From 2e71608c59125f7f016949ceeeffc31b5baac92b Mon Sep 17 00:00:00 2001 From: Rahul Vishwakarma Date: Tue, 26 May 2026 20:05:24 +0530 Subject: [PATCH 1/3] feat(mcp): expose paper-named tools Signed-off-by: Rahul Vishwakarma --- internal/auth/jwt.go | 77 +++++++ internal/auth/jwt_test.go | 45 ++++ internal/hooks/handler.go | 28 +-- internal/mcp/handlers_paper_tools_test.go | 246 ++++++++++++++++++++++ internal/mcp/server.go | 221 +++++++++++++++++++ internal/policy/attenuation.go | 37 ++++ 6 files changed, 627 insertions(+), 27 deletions(-) create mode 100644 internal/mcp/handlers_paper_tools_test.go create mode 100644 internal/policy/attenuation.go diff --git a/internal/auth/jwt.go b/internal/auth/jwt.go index 0c5124a..3beaede 100644 --- a/internal/auth/jwt.go +++ b/internal/auth/jwt.go @@ -117,6 +117,83 @@ func (ti *TokenIssuer) IssueToken( return token.SignedString(ti.signingKey) } +// MintChildToken issues a JWT for a child session bound to a parent +// sublayout. The token's scope is the intersection of the parent +// policy with the sublayout: parent tools (unchanged — sublayouts +// don't broaden tool access), sublayout limits where declared +// (falling back to the parent's limit otherwise), and a fresh +// PolicyDigest computed over the attenuated policy so verification +// can't replay this token against the parent or any other sublayout. +// +// Used by `aflock_delegate` (issue #117) to mirror, in MCP mode, +// the spawn-time binding that hooks-mode performs in +// internal/hooks/handler.go:matchSublayoutForSpawn. +func (ti *TokenIssuer) MintChildToken( + parentPolicy *aflock.Policy, + sub *aflock.Sublayout, + childSessionID string, + agentID string, + identityHash string, + ttl time.Duration, +) (string, error) { + if parentPolicy == nil { + return "", fmt.Errorf("mint child token: nil parent policy") + } + if sub == nil { + return "", fmt.Errorf("mint child token: nil sublayout") + } + child := attenuateChildPolicy(parentPolicy, sub) + return ti.IssueToken(childSessionID, agentID, identityHash, child, ttl) +} + +// attenuateChildPolicy builds an attenuated policy for a sublayout +// child. Tools and grants are carried from the parent (the sublayout +// is a delegation, not a broadening); limits are merged field-by-field +// with the sublayout taking precedence where present. +func attenuateChildPolicy(parent *aflock.Policy, sub *aflock.Sublayout) *aflock.Policy { + child := &aflock.Policy{ + Name: sub.Name, + Version: parent.Version, + Tools: parent.Tools, + } + child.Limits = mergeLimits(parent.Limits, sub.Limits) + return child +} + +// mergeLimits returns a new LimitsPolicy where each field is the +// sublayout's value if set, otherwise the parent's value. Caller is +// responsible for attenuation validation (sub <= parent) before +// calling this — mergeLimits trusts its inputs. +func mergeLimits(parent, sub *aflock.LimitsPolicy) *aflock.LimitsPolicy { + if parent == nil && sub == nil { + return nil + } + out := &aflock.LimitsPolicy{} + pick := func(p, s *aflock.Limit) *aflock.Limit { + if s != nil { + return s + } + return p + } + if parent != nil { + out.MaxSpendUSD = parent.MaxSpendUSD + out.MaxTokensIn = parent.MaxTokensIn + out.MaxTokensOut = parent.MaxTokensOut + out.MaxTurns = parent.MaxTurns + out.MaxWallTimeSeconds = parent.MaxWallTimeSeconds + out.MaxToolCalls = parent.MaxToolCalls + } + if sub != nil { + out.MaxSpendUSD = pick(out.MaxSpendUSD, sub.MaxSpendUSD) + out.MaxTokensIn = pick(out.MaxTokensIn, sub.MaxTokensIn) + out.MaxTokensOut = pick(out.MaxTokensOut, sub.MaxTokensOut) + out.MaxTurns = pick(out.MaxTurns, sub.MaxTurns) + out.MaxWallTimeSeconds = pick(out.MaxWallTimeSeconds, sub.MaxWallTimeSeconds) + out.MaxToolCalls = pick(out.MaxToolCalls, sub.MaxToolCalls) + } + return out +} + // ValidateToken verifies signature, expiry, issuer, and parses claims. func (ti *TokenIssuer) ValidateToken(tokenString string) (*AflockClaims, error) { token, err := jwt.ParseWithClaims(tokenString, &AflockClaims{}, diff --git a/internal/auth/jwt_test.go b/internal/auth/jwt_test.go index bc7826b..852cb57 100644 --- a/internal/auth/jwt_test.go +++ b/internal/auth/jwt_test.go @@ -379,3 +379,48 @@ func TestValidatorRejectsForeignToken(t *testing.T) { _, err = validatorB.ValidateToken(tokenA) assert.Error(t, err) } + +func TestMintChildToken_MergesLimits(t *testing.T) { + issuer, err := NewTokenIssuer() + require.NoError(t, err) + + parent := &aflock.Policy{ + Version: "1.0", + Name: "parent", + Tools: &aflock.ToolsPolicy{Allow: []string{"Read", "Bash"}}, + Limits: &aflock.LimitsPolicy{ + MaxSpendUSD: &aflock.Limit{Value: 10.0, Enforcement: "fail-fast"}, + MaxTokensIn: &aflock.Limit{Value: 500_000, Enforcement: "fail-fast"}, + }, + } + sub := &aflock.Sublayout{ + Name: "research-agent", + Limits: &aflock.LimitsPolicy{ + MaxSpendUSD: &aflock.Limit{Value: 2.0, Enforcement: "fail-fast"}, + // MaxTokensIn not declared — should inherit from parent. + }, + } + + tok, err := issuer.MintChildToken(parent, sub, "child-1", "spiffe://aflock.ai/agent/child", "hash-1", time.Hour) + require.NoError(t, err) + + claims, err := issuer.ValidateTokenForSessionAndPolicy(tok, "child-1", "") + require.NoError(t, err) + + require.NotNil(t, claims.Limits, "child JWT must carry limits") + require.NotNil(t, claims.Limits.MaxSpendUSD, "MaxSpendUSD must be present") + assert.Equal(t, 2.0, claims.Limits.MaxSpendUSD.Value, "child takes sublayout's MaxSpendUSD") + require.NotNil(t, claims.Limits.MaxTokensIn, "MaxTokensIn must inherit from parent") + assert.Equal(t, 500_000.0, claims.Limits.MaxTokensIn.Value, "child inherits parent's MaxTokensIn") +} + +func TestMintChildToken_NilArgs(t *testing.T) { + issuer, err := NewTokenIssuer() + require.NoError(t, err) + + _, err = issuer.MintChildToken(nil, &aflock.Sublayout{Name: "x"}, "child", "agent", "hash", time.Hour) + assert.Error(t, err, "nil parent policy must error") + + _, err = issuer.MintChildToken(&aflock.Policy{}, nil, "child", "agent", "hash", time.Hour) + assert.Error(t, err, "nil sublayout must error") +} diff --git a/internal/hooks/handler.go b/internal/hooks/handler.go index 954a200..b503c5b 100644 --- a/internal/hooks/handler.go +++ b/internal/hooks/handler.go @@ -578,7 +578,7 @@ func (h *Handler) handlePreToolUse(input *aflock.HookInput) error { case hasDecl && matched != nil: // Sublayout limits must already attenuate vs parent's limits; // refusing at spawn time means a bad policy can't slip through. - if violations := attenuationViolations(sessionState.Policy.Limits, matched.Limits); len(violations) > 0 { + if violations := policy.AttenuationViolations(sessionState.Policy.Limits, matched.Limits); len(violations) > 0 { return output.Write(output.PreToolUseDeny(fmt.Sprintf( "[aflock] BLOCKED: sublayout %q violates parent attenuation: %s", matched.Name, strings.Join(violations, "; ")))) @@ -1570,32 +1570,6 @@ func matchSublayoutForSpawn(toolName string, toolInput json.RawMessage, sublayou return nil, true } -// attenuationViolations checks that sublayout limits are <= parent limits. -// Mirrors verify.verifyAttenuation but lives here so spawn-time enforcement -// (issue #26 gap 4) doesn't pull in the verify package's dependencies. -// Empty result means the sublayout is validly attenuated. -func attenuationViolations(parent, sub *aflock.LimitsPolicy) []string { - if sub == nil || parent == nil { - return nil - } - var violations []string - check := func(name string, p, s *aflock.Limit) { - if p == nil || s == nil { - return - } - if s.Value > p.Value { - violations = append(violations, fmt.Sprintf("%s: sublayout %.2f > parent %.2f", name, s.Value, p.Value)) - } - } - check("maxSpendUSD", parent.MaxSpendUSD, sub.MaxSpendUSD) - check("maxTokensIn", parent.MaxTokensIn, sub.MaxTokensIn) - check("maxTokensOut", parent.MaxTokensOut, sub.MaxTokensOut) - check("maxTurns", parent.MaxTurns, sub.MaxTurns) - check("maxWallTimeSeconds", parent.MaxWallTimeSeconds, sub.MaxWallTimeSeconds) - check("maxToolCalls", parent.MaxToolCalls, sub.MaxToolCalls) - return violations -} - // attenuateLimits computes effective limits for a child session. // For each limit field: child effective = min(child policy limit, parent remaining). // If a parent has exhausted its budget, the child gets 0. diff --git a/internal/mcp/handlers_paper_tools_test.go b/internal/mcp/handlers_paper_tools_test.go new file mode 100644 index 0000000..fe5ca41 --- /dev/null +++ b/internal/mcp/handlers_paper_tools_test.go @@ -0,0 +1,246 @@ +package mcp + +import ( + "context" + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "encoding/json" + "strings" + "testing" + "time" + + "github.com/mark3labs/mcp-go/mcp" + + "github.com/aflock-ai/aflock/internal/auth" + "github.com/aflock-ai/aflock/pkg/aflock" +) + +// policyWithLimits returns a policy with all six limit categories declared +// so handleCheckLimits has something to surface for each. +func policyWithLimits() *aflock.Policy { + return &aflock.Policy{ + Version: "1.0", + Name: "check-limits-fixture", + Tools: &aflock.ToolsPolicy{Allow: []string{"Read", "Bash"}}, + Limits: &aflock.LimitsPolicy{ + MaxSpendUSD: &aflock.Limit{Value: 10.0, Enforcement: "fail-fast"}, + MaxTokensIn: &aflock.Limit{Value: 500_000, Enforcement: "fail-fast"}, + MaxTokensOut: &aflock.Limit{Value: 200_000, Enforcement: "post-hoc"}, + MaxTurns: &aflock.Limit{Value: 50, Enforcement: "post-hoc"}, + MaxToolCalls: &aflock.Limit{Value: 100, Enforcement: "fail-fast"}, + MaxWallTimeSeconds: &aflock.Limit{Value: 3600, Enforcement: "fail-fast"}, + }, + } +} + +// TestHandleCheckLimits_SurfacesAllDeclared confirms every declared +// limit appears in the response with the {value, used, remaining, +// enforcement} shape the paper-spec consumers will rely on. +func TestHandleCheckLimits_SurfacesAllDeclared(t *testing.T) { + s := newTestServerWithPolicy(t, policyWithLimits()) + + // Seed some consumption so used/remaining aren't both zero. + sess, err := s.stateManager.Load(s.sessionID) + if err != nil || sess == nil { + t.Fatalf("load session: %v", err) + } + sess.Metrics.CostUSD = 2.5 + sess.Metrics.TokensIn = 12_000 + sess.Metrics.TokensOut = 3_000 + sess.Metrics.Turns = 7 + sess.Metrics.ToolCalls = 9 + if err := s.stateManager.Save(sess); err != nil { + t.Fatalf("save session: %v", err) + } + + result, err := s.handleCheckLimits(context.Background(), newTestRequest(nil)) + if err != nil { + t.Fatalf("handleCheckLimits: %v", err) + } + + var parsed struct { + Limits map[string]map[string]any `json:"limits"` + } + if err := json.Unmarshal([]byte(result.Content[0].(mcp.TextContent).Text), &parsed); err != nil { + t.Fatalf("parse response: %v", err) + } + + for _, name := range []string{ + "maxSpendUSD", "maxTokensIn", "maxTokensOut", + "maxTurns", "maxToolCalls", "maxWallTimeSeconds", + } { + entry, ok := parsed.Limits[name] + if !ok { + t.Errorf("missing limit %q in response", name) + continue + } + for _, field := range []string{"value", "used", "remaining", "enforcement"} { + if _, ok := entry[field]; !ok { + t.Errorf("limit %q missing %q field", name, field) + } + } + } + + // Spot-check the numbers: cost limit 10.0, used 2.5 → remaining 7.5. + spend := parsed.Limits["maxSpendUSD"] + if spend["used"].(float64) != 2.5 || spend["remaining"].(float64) != 7.5 { + t.Errorf("maxSpendUSD used/remaining wrong: %v", spend) + } +} + +// TestHandleCheckLimits_NoPolicy returns empty `limits` rather than erroring, +// so callers can poll the endpoint before policy loads. +func TestHandleCheckLimits_NoPolicy(t *testing.T) { + s := newTestServer(t) + result, err := s.handleCheckLimits(context.Background(), newTestRequest(nil)) + if err != nil { + t.Fatalf("handleCheckLimits: %v", err) + } + var parsed struct { + Limits map[string]any `json:"limits"` + } + if err := json.Unmarshal([]byte(result.Content[0].(mcp.TextContent).Text), &parsed); err != nil { + t.Fatalf("parse response: %v", err) + } + if len(parsed.Limits) != 0 { + t.Errorf("expected empty limits with no policy, got %v", parsed.Limits) + } +} + +// policyWithSublayout returns a policy declaring one valid sublayout +// (`research-agent`) so handleDelegate has something to match. +func policyWithSublayout() *aflock.Policy { + return &aflock.Policy{ + Version: "1.0", + Name: "delegate-fixture", + Tools: &aflock.ToolsPolicy{Allow: []string{"Read", "Bash"}}, + Limits: &aflock.LimitsPolicy{ + MaxSpendUSD: &aflock.Limit{Value: 10.0, Enforcement: "fail-fast"}, + MaxTokensIn: &aflock.Limit{Value: 500_000, Enforcement: "fail-fast"}, + }, + Sublayouts: []aflock.Sublayout{{ + Name: "research-agent", + Policy: "./policies/research.aflock", + Limits: &aflock.LimitsPolicy{ + MaxSpendUSD: &aflock.Limit{Value: 2.0, Enforcement: "fail-fast"}, + }, + }}, + } +} + +// TestHandleDelegate_MissingSublayoutName rejects empty input clearly. +func TestHandleDelegate_MissingSublayoutName(t *testing.T) { + s := newTestServerWithPolicy(t, policyWithSublayout()) + result, err := s.handleDelegate(context.Background(), newTestRequest(nil)) + if err != nil { + t.Fatalf("handleDelegate: %v", err) + } + if !result.IsError || !strings.Contains(result.Content[0].(mcp.TextContent).Text, "sublayout_name is required") { + t.Errorf("expected 'sublayout_name is required' error, got: %+v", result) + } +} + +// TestHandleDelegate_UnknownSublayout refuses names not in policy. +func TestHandleDelegate_UnknownSublayout(t *testing.T) { + s := newTestServerWithPolicy(t, policyWithSublayout()) + result, err := s.handleDelegate(context.Background(), + newTestRequest(map[string]any{"sublayout_name": "nope"})) + if err != nil { + t.Fatalf("handleDelegate: %v", err) + } + if !result.IsError || !strings.Contains(result.Content[0].(mcp.TextContent).Text, "not declared in policy") { + t.Errorf("expected 'not declared' error, got: %+v", result) + } +} + +// TestHandleDelegate_AttenuationViolation rejects sublayouts whose +// limits exceed the parent's. +func TestHandleDelegate_AttenuationViolation(t *testing.T) { + pol := policyWithSublayout() + // Bump the sublayout above the parent. + pol.Sublayouts[0].Limits.MaxSpendUSD = &aflock.Limit{Value: 99.0, Enforcement: "fail-fast"} + s := newTestServerWithPolicy(t, pol) + + result, err := s.handleDelegate(context.Background(), + newTestRequest(map[string]any{"sublayout_name": "research-agent"})) + if err != nil { + t.Fatalf("handleDelegate: %v", err) + } + if !result.IsError || !strings.Contains(result.Content[0].(mcp.TextContent).Text, "attenuation") { + t.Errorf("expected 'attenuation' error, got: %+v", result) + } +} + +// TestHandleDelegate_HappyPath validates the full success flow: +// returns ok=true, includes the sublayout name, mints a child JWT, and +// the JWT's claims carry the attenuated limit. +func TestHandleDelegate_HappyPath(t *testing.T) { + s := newTestServerWithPolicy(t, policyWithSublayout()) + key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + if err != nil { + t.Fatalf("generate key: %v", err) + } + s.tokenIssuer = auth.NewTokenIssuerFromSigner(key, "test-signer") + + result, err := s.handleDelegate(context.Background(), + newTestRequest(map[string]any{"sublayout_name": "research-agent"})) + if err != nil { + t.Fatalf("handleDelegate: %v", err) + } + if result.IsError { + t.Fatalf("unexpected error: %s", result.Content[0].(mcp.TextContent).Text) + } + + var parsed struct { + OK bool `json:"ok"` + Sublayout string `json:"sublayout"` + ChildSessionID string `json:"child_session_id"` + ChildJWT string `json:"child_jwt"` + } + if err := json.Unmarshal([]byte(result.Content[0].(mcp.TextContent).Text), &parsed); err != nil { + t.Fatalf("parse response: %v", err) + } + if !parsed.OK || parsed.Sublayout != "research-agent" { + t.Errorf("unexpected response shape: %+v", parsed) + } + if parsed.ChildSessionID == "" || parsed.ChildJWT == "" { + t.Errorf("child_session_id / child_jwt must be populated: %+v", parsed) + } + + // Validate the minted JWT carries the sublayout's attenuated limit. + // Pass empty currentPolicyDigest because the child JWT is bound to the + // attenuated child policy, not the parent's digest. + claims, err := s.tokenIssuer.ValidateTokenForSessionAndPolicy(parsed.ChildJWT, parsed.ChildSessionID, "") + if err != nil { + t.Fatalf("validate child JWT: %v", err) + } + if claims.Limits == nil || claims.Limits.MaxSpendUSD == nil || claims.Limits.MaxSpendUSD.Value != 2.0 { + t.Errorf("child JWT must carry sublayout MaxSpendUSD=2.0, got: %+v", claims.Limits) + } + // And it inherits the parent's TokensIn limit unchanged. + if claims.Limits.MaxTokensIn == nil || claims.Limits.MaxTokensIn.Value != 500_000 { + t.Errorf("child JWT must inherit parent MaxTokensIn=500000, got: %+v", claims.Limits.MaxTokensIn) + } +} + +// TestRegisterTools_PaperNamesPresent confirms the four paper-named +// tools are registered (alongside the legacy names) so a server probe +// won't hit `tool not found` on any of them. +func TestRegisterTools_PaperNamesPresent(t *testing.T) { + s := NewServer() + listed := s.mcpServer.ListTools() + for _, name := range []string{ + "aflock_authorize", "aflock_attest", + "aflock_check_limits", "aflock_delegate", + // legacy names retained for backward compatibility: + "check_tool", "sign_attestation", + } { + if _, ok := listed[name]; !ok { + t.Errorf("expected tool %q registered, missing", name) + } + } + // Compile-time anchor — keep the import live even if the body + // changes shape later. + _ = time.Second +} diff --git a/internal/mcp/server.go b/internal/mcp/server.go index a9270c1..f31199a 100644 --- a/internal/mcp/server.go +++ b/internal/mcp/server.go @@ -200,6 +200,51 @@ func (s *Server) registerTools() { ), s.handleSignAttestation, ) + + // Paper-named tools (issue #117, paper §3.3). aflock_authorize and + // aflock_attest are aliases of check_tool and sign_attestation so + // readers of the paper can probe the server with the names the paper + // uses. check_tool and sign_attestation remain registered for + // backward compatibility. + + // aflock_authorize - Paper §3.3 alias of check_tool. + s.mcpServer.AddTool( + mcp.NewTool("aflock_authorize", + mcp.WithDescription("Request authorization for an action (paper §3.3). Alias of check_tool."), + mcp.WithString("tool_name", mcp.Required(), mcp.Description("Name of the tool to check")), + mcp.WithObject("tool_input", mcp.Description("Tool input parameters")), + ), + s.handleCheckTool, + ) + + // aflock_attest - Paper §3.3 alias of sign_attestation. + s.mcpServer.AddTool( + mcp.NewTool("aflock_attest", + mcp.WithDescription("Record an action — server signs attestation (paper §3.3). Alias of sign_attestation."), + mcp.WithString("predicate_type", mcp.Required(), mcp.Description("Predicate type URI")), + mcp.WithObject("predicate", mcp.Required(), mcp.Description("Predicate data to attest")), + mcp.WithObject("subject", mcp.Description("Subject to bind attestation to (name and digest)")), + ), + s.handleSignAttestation, + ) + + // aflock_check_limits - Paper §3.3, returns remaining budget per limit. + s.mcpServer.AddTool( + mcp.NewTool("aflock_check_limits", + mcp.WithDescription("Query remaining budget against each declared policy limit (paper §3.3)."), + ), + s.handleCheckLimits, + ) + + // aflock_delegate - Paper §3.3, creates a sublayout binding for a sub-agent. + s.mcpServer.AddTool( + mcp.NewTool("aflock_delegate", + mcp.WithDescription("Create a sublayout binding for a sub-agent (paper §3.3). Validates the named sublayout against parent attenuation, writes a propagation record, and returns an attenuated child JWT."), + mcp.WithString("sublayout_name", mcp.Required(), mcp.Description("Name of a sublayout declared in the parent policy")), + mcp.WithString("child_session_id", mcp.Description("Optional pre-bound child session ID; one is generated if omitted")), + ), + s.handleDelegate, + ) } // Serve starts the MCP server on stdio. @@ -1634,6 +1679,182 @@ func (s *Server) handleSignAttestation(ctx context.Context, request mcp.CallTool return mcp.NewToolResultText(string(data)), nil } +// handleCheckLimits returns the configured resource limits alongside +// their current consumption and remaining budget. Paper §3.3 names +// this `aflock_check_limits`. Read-only: no JWT required, since this +// just exposes data already returned by `get_session` and the loaded +// policy. Omits any limit category the policy hasn't declared. +func (s *Server) handleCheckLimits(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { + if s.policy == nil || s.policy.Limits == nil { + return mcp.NewToolResultText(`{"limits": {}}`), nil + } + + sessionState, err := s.stateManager.Load(s.sessionID) + if err != nil { + return nil, fmt.Errorf("load session: %w", err) + } + + var ( + costUSD float64 + tokensIn int64 + tokensOut int64 + turns int + toolCalls int + ) + startedAt := time.Now() + if sessionState != nil { + startedAt = sessionState.StartedAt + if sessionState.Metrics != nil { + costUSD = sessionState.Metrics.CostUSD + tokensIn = sessionState.Metrics.TokensIn + tokensOut = sessionState.Metrics.TokensOut + turns = sessionState.Metrics.Turns + toolCalls = sessionState.Metrics.ToolCalls + } + } + + limits := map[string]any{} + addFloat := func(name string, lim *aflock.Limit, used float64) { + if lim == nil { + return + } + remaining := lim.Value - used + if remaining < 0 { + remaining = 0 + } + limits[name] = map[string]any{ + "value": lim.Value, + "used": used, + "remaining": remaining, + "enforcement": lim.Enforcement, + } + } + addInt := func(name string, lim *aflock.Limit, used int64) { + if lim == nil { + return + } + remaining := int64(lim.Value) - used + if remaining < 0 { + remaining = 0 + } + limits[name] = map[string]any{ + "value": int64(lim.Value), + "used": used, + "remaining": remaining, + "enforcement": lim.Enforcement, + } + } + + pol := s.policy.Limits + addFloat("maxSpendUSD", pol.MaxSpendUSD, costUSD) + addInt("maxTokensIn", pol.MaxTokensIn, tokensIn) + addInt("maxTokensOut", pol.MaxTokensOut, tokensOut) + addInt("maxTurns", pol.MaxTurns, int64(turns)) + addInt("maxToolCalls", pol.MaxToolCalls, int64(toolCalls)) + if pol.MaxWallTimeSeconds != nil { + elapsed := time.Since(startedAt).Seconds() + if elapsed < 0 { + elapsed = 0 + } + addFloat("maxWallTimeSeconds", pol.MaxWallTimeSeconds, elapsed) + } + + result := map[string]any{"limits": limits} + data, _ := json.MarshalIndent(result, "", " ") + return mcp.NewToolResultText(string(data)), nil +} + +// handleDelegate creates a sublayout binding for a sub-agent. Paper +// §3.3 names this `aflock_delegate`. JWT-gated (parent agent must be +// authorized). Validates the named sublayout exists, checks +// attenuation, writes a propagation record, and mints an attenuated +// child JWT. +// +// This mirrors, in MCP mode, the spawn-time binding that +// internal/hooks/handler.go performs at PreToolUse when Claude Code's +// Task/Agent tool is invoked. +func (s *Server) handleDelegate(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { + claims, err := s.validateJWT(request) + if err != nil { + return mcp.NewToolResultError(fmt.Sprintf("Authorization denied: %v", err)), nil + } + if claims == nil && s.authActive.Load() { + return mcp.NewToolResultError("Authorization denied: aflock_delegate requires a valid JWT once auth is active"), nil + } + + if s.policy == nil { + return mcp.NewToolResultError("aflock_delegate: no policy loaded"), nil + } + if err := s.errPolicyExpired(); err != nil { + return mcp.NewToolResultError(err.Error()), nil + } + + sublayoutName := request.GetString("sublayout_name", "") + if sublayoutName == "" { + return mcp.NewToolResultError("sublayout_name is required"), nil + } + + var matched *aflock.Sublayout + for i := range s.policy.Sublayouts { + if s.policy.Sublayouts[i].Name == sublayoutName { + matched = &s.policy.Sublayouts[i] + break + } + } + if matched == nil { + return mcp.NewToolResultError(fmt.Sprintf("aflock_delegate: sublayout %q not declared in policy", sublayoutName)), nil + } + + if violations := policy.AttenuationViolations(s.policy.Limits, matched.Limits); len(violations) > 0 { + return mcp.NewToolResultError(fmt.Sprintf("aflock_delegate: sublayout %q violates parent attenuation: %s", + sublayoutName, strings.Join(violations, "; "))), nil + } + + parentState, err := s.stateManager.Load(s.sessionID) + if err != nil { + return nil, fmt.Errorf("load parent session: %w", err) + } + if parentState == nil { + return mcp.NewToolResultError("aflock_delegate: no parent session state; call get_token / start a session first"), nil + } + + if err := s.stateManager.WritePropagationForSublayout(parentState, matched); err != nil { + return mcp.NewToolResultError(fmt.Sprintf("aflock_delegate: write propagation: %v", err)), nil + } + + childSessionID := request.GetString("child_session_id", "") + if childSessionID == "" { + childSessionID = fmt.Sprintf("delegate-%s", uuid.New().String()) + } + + var childJWT string + if s.tokenIssuer != nil { + agentID := "" + identityHash := "" + if s.agentIdentity != nil { + if spiffeID, sErr := s.agentIdentity.ToSPIFFEID("aflock.ai"); sErr == nil { + agentID = spiffeID.String() + } + identityHash = s.agentIdentity.IdentityHash + } + childJWT, err = s.tokenIssuer.MintChildToken(s.policy, matched, childSessionID, agentID, identityHash, time.Hour) + if err != nil { + return mcp.NewToolResultError(fmt.Sprintf("aflock_delegate: mint child JWT: %v", err)), nil + } + } + + result := map[string]any{ + "ok": true, + "sublayout": matched.Name, + "parent_session": s.sessionID, + "child_session_id": childSessionID, + "child_jwt": childJWT, + "limits": matched.Limits, + } + data, _ := json.MarshalIndent(result, "", " ") + return mcp.NewToolResultText(string(data)), nil +} + // computePredicateDigest computes SHA256 of predicate data. func computePredicateDigest(predicate interface{}) string { data, _ := json.Marshal(predicate) diff --git a/internal/policy/attenuation.go b/internal/policy/attenuation.go new file mode 100644 index 0000000..01f3a90 --- /dev/null +++ b/internal/policy/attenuation.go @@ -0,0 +1,37 @@ +package policy + +import ( + "fmt" + + "github.com/aflock-ai/aflock/pkg/aflock" +) + +// AttenuationViolations checks that sublayout limits are <= parent limits. +// Empty result means the sublayout is validly attenuated. Sublayouts may +// declare fewer limits than the parent; only fields set on both sides are +// compared. A nil parent or sub means there's nothing to compare against. +// +// Used at spawn time (hooks-mode PreToolUse, MCP-mode aflock_delegate) +// and at verify time. Sharing one rule keeps the two enforcement paths +// from drifting. +func AttenuationViolations(parent, sub *aflock.LimitsPolicy) []string { + if sub == nil || parent == nil { + return nil + } + var violations []string + check := func(name string, p, s *aflock.Limit) { + if p == nil || s == nil { + return + } + if s.Value > p.Value { + violations = append(violations, fmt.Sprintf("%s: sublayout %.2f > parent %.2f", name, s.Value, p.Value)) + } + } + check("maxSpendUSD", parent.MaxSpendUSD, sub.MaxSpendUSD) + check("maxTokensIn", parent.MaxTokensIn, sub.MaxTokensIn) + check("maxTokensOut", parent.MaxTokensOut, sub.MaxTokensOut) + check("maxTurns", parent.MaxTurns, sub.MaxTurns) + check("maxWallTimeSeconds", parent.MaxWallTimeSeconds, sub.MaxWallTimeSeconds) + check("maxToolCalls", parent.MaxToolCalls, sub.MaxToolCalls) + return violations +} From a8c545848339d4d755258bc429f0183e4f61e1eb Mon Sep 17 00:00:00 2001 From: Rahul Vishwakarma Date: Tue, 26 May 2026 20:34:50 +0530 Subject: [PATCH 2/3] fix(mcp): tighten aflock_delegate per copilot review Signed-off-by: Rahul Vishwakarma --- internal/mcp/handlers_paper_tools_test.go | 80 +++++++++++++++++++---- internal/mcp/server.go | 73 +++++++++++++++------ internal/verify/sublayout_test.go | 2 +- internal/verify/verifier.go | 34 ++-------- 4 files changed, 126 insertions(+), 63 deletions(-) diff --git a/internal/mcp/handlers_paper_tools_test.go b/internal/mcp/handlers_paper_tools_test.go index fe5ca41..12dab35 100644 --- a/internal/mcp/handlers_paper_tools_test.go +++ b/internal/mcp/handlers_paper_tools_test.go @@ -129,10 +129,51 @@ func policyWithSublayout() *aflock.Policy { } } +// authedDelegateServer wires up a server with the sublayout fixture and +// pre-issues a parent JWT. Returns the server and the token string callers +// pass back through `_token` so handleDelegate's JWT gate is satisfied. +func authedDelegateServer(t *testing.T, pol *aflock.Policy) (*Server, string) { + t.Helper() + s := newTestServerWithPolicy(t, pol) + key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + if err != nil { + t.Fatalf("generate key: %v", err) + } + s.tokenIssuer = auth.NewTokenIssuerFromSigner(key, "test-signer") + tok, err := s.tokenIssuer.IssueToken(s.sessionID, "test-agent", "test-hash", pol, time.Hour) + if err != nil { + t.Fatalf("issue parent token: %v", err) + } + return s, tok +} + +// TestHandleDelegate_RequiresJWT — handler must refuse calls without a +// JWT regardless of authActive state, since delegate mints child tokens. +func TestHandleDelegate_RequiresJWT(t *testing.T) { + s := newTestServerWithPolicy(t, policyWithSublayout()) + key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + if err != nil { + t.Fatalf("generate key: %v", err) + } + s.tokenIssuer = auth.NewTokenIssuerFromSigner(key, "test-signer") + // authActive deliberately left false — graceful adoption must NOT + // apply to aflock_delegate. + + result, err := s.handleDelegate(context.Background(), + newTestRequest(map[string]any{"sublayout_name": "research-agent"})) + if err != nil { + t.Fatalf("handleDelegate: %v", err) + } + if !result.IsError || !strings.Contains(result.Content[0].(mcp.TextContent).Text, "requires a valid JWT") { + t.Errorf("expected 'requires a valid JWT' error, got: %+v", result) + } +} + // TestHandleDelegate_MissingSublayoutName rejects empty input clearly. func TestHandleDelegate_MissingSublayoutName(t *testing.T) { - s := newTestServerWithPolicy(t, policyWithSublayout()) - result, err := s.handleDelegate(context.Background(), newTestRequest(nil)) + s, tok := authedDelegateServer(t, policyWithSublayout()) + result, err := s.handleDelegate(context.Background(), + newTestRequest(map[string]any{"_token": tok})) if err != nil { t.Fatalf("handleDelegate: %v", err) } @@ -143,9 +184,9 @@ func TestHandleDelegate_MissingSublayoutName(t *testing.T) { // TestHandleDelegate_UnknownSublayout refuses names not in policy. func TestHandleDelegate_UnknownSublayout(t *testing.T) { - s := newTestServerWithPolicy(t, policyWithSublayout()) + s, tok := authedDelegateServer(t, policyWithSublayout()) result, err := s.handleDelegate(context.Background(), - newTestRequest(map[string]any{"sublayout_name": "nope"})) + newTestRequest(map[string]any{"_token": tok, "sublayout_name": "nope"})) if err != nil { t.Fatalf("handleDelegate: %v", err) } @@ -160,10 +201,10 @@ func TestHandleDelegate_AttenuationViolation(t *testing.T) { pol := policyWithSublayout() // Bump the sublayout above the parent. pol.Sublayouts[0].Limits.MaxSpendUSD = &aflock.Limit{Value: 99.0, Enforcement: "fail-fast"} - s := newTestServerWithPolicy(t, pol) + s, tok := authedDelegateServer(t, pol) result, err := s.handleDelegate(context.Background(), - newTestRequest(map[string]any{"sublayout_name": "research-agent"})) + newTestRequest(map[string]any{"_token": tok, "sublayout_name": "research-agent"})) if err != nil { t.Fatalf("handleDelegate: %v", err) } @@ -172,19 +213,32 @@ func TestHandleDelegate_AttenuationViolation(t *testing.T) { } } +// TestHandleDelegate_RejectsBadChildSessionID — caller-supplied IDs +// with disallowed characters must be rejected before any side effect. +func TestHandleDelegate_RejectsBadChildSessionID(t *testing.T) { + s, tok := authedDelegateServer(t, policyWithSublayout()) + result, err := s.handleDelegate(context.Background(), + newTestRequest(map[string]any{ + "_token": tok, + "sublayout_name": "research-agent", + "child_session_id": "../../etc/passwd", + })) + if err != nil { + t.Fatalf("handleDelegate: %v", err) + } + if !result.IsError || !strings.Contains(result.Content[0].(mcp.TextContent).Text, "child_session_id must be") { + t.Errorf("expected child_session_id format error, got: %+v", result) + } +} + // TestHandleDelegate_HappyPath validates the full success flow: // returns ok=true, includes the sublayout name, mints a child JWT, and // the JWT's claims carry the attenuated limit. func TestHandleDelegate_HappyPath(t *testing.T) { - s := newTestServerWithPolicy(t, policyWithSublayout()) - key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) - if err != nil { - t.Fatalf("generate key: %v", err) - } - s.tokenIssuer = auth.NewTokenIssuerFromSigner(key, "test-signer") + s, tok := authedDelegateServer(t, policyWithSublayout()) result, err := s.handleDelegate(context.Background(), - newTestRequest(map[string]any{"sublayout_name": "research-agent"})) + newTestRequest(map[string]any{"_token": tok, "sublayout_name": "research-agent"})) if err != nil { t.Fatalf("handleDelegate: %v", err) } diff --git a/internal/mcp/server.go b/internal/mcp/server.go index f31199a..04bee65 100644 --- a/internal/mcp/server.go +++ b/internal/mcp/server.go @@ -1774,12 +1774,18 @@ func (s *Server) handleCheckLimits(ctx context.Context, request mcp.CallToolRequ // internal/hooks/handler.go performs at PreToolUse when Claude Code's // Task/Agent tool is invoked. func (s *Server) handleDelegate(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { + // aflock_delegate is sensitive (issues child JWTs + writes propagation), + // so we require a validated JWT unconditionally rather than honoring + // the graceful-adoption fallback used by less-privileged tools. claims, err := s.validateJWT(request) if err != nil { return mcp.NewToolResultError(fmt.Sprintf("Authorization denied: %v", err)), nil } - if claims == nil && s.authActive.Load() { - return mcp.NewToolResultError("Authorization denied: aflock_delegate requires a valid JWT once auth is active"), nil + if claims == nil { + return mcp.NewToolResultError("Authorization denied: aflock_delegate requires a valid JWT"), nil + } + if s.tokenIssuer == nil { + return mcp.NewToolResultError("aflock_delegate: token issuer not initialized"), nil } if s.policy == nil { @@ -1818,29 +1824,36 @@ func (s *Server) handleDelegate(ctx context.Context, request mcp.CallToolRequest return mcp.NewToolResultError("aflock_delegate: no parent session state; call get_token / start a session first"), nil } - if err := s.stateManager.WritePropagationForSublayout(parentState, matched); err != nil { - return mcp.NewToolResultError(fmt.Sprintf("aflock_delegate: write propagation: %v", err)), nil - } - + // Caller-supplied child_session_id must be a safe identifier the + // state manager and downstream tooling can use as a filename + // component. We allow alphanumerics, hyphen, and underscore; + // anything else is rejected so a malformed id can't slip through + // to a JWT and then fail at state-load time. childSessionID := request.GetString("child_session_id", "") if childSessionID == "" { childSessionID = fmt.Sprintf("delegate-%s", uuid.New().String()) + } else if !validSessionID(childSessionID) { + return mcp.NewToolResultError( + "aflock_delegate: child_session_id must be 1-128 chars of [a-zA-Z0-9_-]"), nil } - var childJWT string - if s.tokenIssuer != nil { - agentID := "" - identityHash := "" - if s.agentIdentity != nil { - if spiffeID, sErr := s.agentIdentity.ToSPIFFEID("aflock.ai"); sErr == nil { - agentID = spiffeID.String() - } - identityHash = s.agentIdentity.IdentityHash - } - childJWT, err = s.tokenIssuer.MintChildToken(s.policy, matched, childSessionID, agentID, identityHash, time.Hour) - if err != nil { - return mcp.NewToolResultError(fmt.Sprintf("aflock_delegate: mint child JWT: %v", err)), nil + // Mint the JWT first. If minting fails, no propagation record is + // written — failed delegations have no observable side effects. + agentID := "" + identityHash := "" + if s.agentIdentity != nil { + if spiffeID, sErr := s.agentIdentity.ToSPIFFEID("aflock.ai"); sErr == nil { + agentID = spiffeID.String() } + identityHash = s.agentIdentity.IdentityHash + } + childJWT, err := s.tokenIssuer.MintChildToken(s.policy, matched, childSessionID, agentID, identityHash, time.Hour) + if err != nil { + return mcp.NewToolResultError(fmt.Sprintf("aflock_delegate: mint child JWT: %v", err)), nil + } + + if err := s.stateManager.WritePropagationForSublayout(parentState, matched); err != nil { + return mcp.NewToolResultError(fmt.Sprintf("aflock_delegate: write propagation: %v", err)), nil } result := map[string]any{ @@ -1855,6 +1868,28 @@ func (s *Server) handleDelegate(ctx context.Context, request mcp.CallToolRequest return mcp.NewToolResultText(string(data)), nil } +// validSessionID enforces the safe-identifier rule for caller-supplied +// child_session_id values: 1-128 chars, restricted to [a-zA-Z0-9_-]. +// Keeps malformed IDs from reaching the state manager (where they'd +// become filename components) and from being baked into a JWT that +// downstream tooling can't load. +func validSessionID(s string) bool { + if len(s) == 0 || len(s) > 128 { + return false + } + for _, r := range s { + switch { + case r >= 'a' && r <= 'z': + case r >= 'A' && r <= 'Z': + case r >= '0' && r <= '9': + case r == '-' || r == '_': + default: + return false + } + } + return true +} + // computePredicateDigest computes SHA256 of predicate data. func computePredicateDigest(predicate interface{}) string { data, _ := json.Marshal(predicate) diff --git a/internal/verify/sublayout_test.go b/internal/verify/sublayout_test.go index 687f569..e25d2a3 100644 --- a/internal/verify/sublayout_test.go +++ b/internal/verify/sublayout_test.go @@ -56,7 +56,7 @@ func TestVerifyAttenuation_Violation_SpendExceeds(t *testing.T) { if len(violations) != 1 { t.Fatalf("Expected 1 violation, got %d: %v", len(violations), violations) } - if violations[0] != "maxSpendUSD: child 10.00 > parent 5.00" { + if violations[0] != "maxSpendUSD: sublayout 10.00 > parent 5.00" { t.Errorf("Unexpected violation: %s", violations[0]) } } diff --git a/internal/verify/verifier.go b/internal/verify/verifier.go index 650d447..44d5014 100644 --- a/internal/verify/verifier.go +++ b/internal/verify/verifier.go @@ -1657,37 +1657,11 @@ func matchesSublayout(child *aflock.SessionState, sub *aflock.Sublayout) bool { } // verifyAttenuation checks that sub-agent limits are ≤ parent limits. -// This prevents privilege escalation: a sub-agent cannot have higher limits than its parent. -// Returns a list of violations. Empty list = attenuation is valid. +// Thin wrapper over policy.AttenuationViolations so verify-time and +// spawn-time (hooks PreToolUse, MCP aflock_delegate) enforce the same +// rule and can't drift apart. func verifyAttenuation(parent, child *aflock.LimitsPolicy) []string { - if child == nil { - return nil // No child limits = inherits parent (always valid) - } - if parent == nil { - return nil // No parent limits = no constraints to violate - } - - var violations []string - - checkLimit := func(name string, parentLimit, childLimit *aflock.Limit) { - if childLimit == nil || parentLimit == nil { - return // No limit set on one side = no violation - } - if childLimit.Value > parentLimit.Value { - violations = append(violations, fmt.Sprintf( - "%s: child %.2f > parent %.2f", - name, childLimit.Value, parentLimit.Value)) - } - } - - checkLimit("maxSpendUSD", parent.MaxSpendUSD, child.MaxSpendUSD) - checkLimit("maxTokensIn", parent.MaxTokensIn, child.MaxTokensIn) - checkLimit("maxTokensOut", parent.MaxTokensOut, child.MaxTokensOut) - checkLimit("maxTurns", parent.MaxTurns, child.MaxTurns) - checkLimit("maxWallTimeSeconds", parent.MaxWallTimeSeconds, child.MaxWallTimeSeconds) - checkLimit("maxToolCalls", parent.MaxToolCalls, child.MaxToolCalls) - - return violations + return policy.AttenuationViolations(parent, child) } // loadAttestationStatements reads the session's *.intoto.json envelopes, From e29791e5258d2faa859b8e83708a199c4f41309c Mon Sep 17 00:00:00 2001 From: Rahul Vishwakarma Date: Tue, 26 May 2026 22:21:15 +0530 Subject: [PATCH 3/3] fix(mcp): scope-check aflock_delegate and alias-aware sign_attestation Signed-off-by: Rahul Vishwakarma --- internal/mcp/handlers_paper_tools_test.go | 49 +++++++++++++++++++---- internal/mcp/server.go | 31 ++++++++++++-- 2 files changed, 69 insertions(+), 11 deletions(-) diff --git a/internal/mcp/handlers_paper_tools_test.go b/internal/mcp/handlers_paper_tools_test.go index 12dab35..1cd3f0b 100644 --- a/internal/mcp/handlers_paper_tools_test.go +++ b/internal/mcp/handlers_paper_tools_test.go @@ -109,12 +109,15 @@ func TestHandleCheckLimits_NoPolicy(t *testing.T) { } // policyWithSublayout returns a policy declaring one valid sublayout -// (`research-agent`) so handleDelegate has something to match. +// (`research-agent`) so handleDelegate has something to match. The +// allowlist includes `aflock_delegate` so the new tool-scope check +// passes for the happy-path tests; a separate fixture below omits it +// to exercise the scope rejection. func policyWithSublayout() *aflock.Policy { return &aflock.Policy{ Version: "1.0", Name: "delegate-fixture", - Tools: &aflock.ToolsPolicy{Allow: []string{"Read", "Bash"}}, + Tools: &aflock.ToolsPolicy{Allow: []string{"Read", "Bash", "aflock_delegate"}}, Limits: &aflock.LimitsPolicy{ MaxSpendUSD: &aflock.Limit{Value: 10.0, Enforcement: "fail-fast"}, MaxTokensIn: &aflock.Limit{Value: 500_000, Enforcement: "fail-fast"}, @@ -129,6 +132,15 @@ func policyWithSublayout() *aflock.Policy { } } +// delegateRequest mirrors a real MCP server's CallToolRequest by +// stamping the invoked tool name into Params.Name — the scope check +// in handleDelegate reads it. Defaults to "aflock_delegate". +func delegateRequest(args map[string]any) mcp.CallToolRequest { + r := newTestRequest(args) + r.Params.Name = "aflock_delegate" + return r +} + // authedDelegateServer wires up a server with the sublayout fixture and // pre-issues a parent JWT. Returns the server and the token string callers // pass back through `_token` so handleDelegate's JWT gate is satisfied. @@ -160,7 +172,7 @@ func TestHandleDelegate_RequiresJWT(t *testing.T) { // apply to aflock_delegate. result, err := s.handleDelegate(context.Background(), - newTestRequest(map[string]any{"sublayout_name": "research-agent"})) + delegateRequest(map[string]any{"sublayout_name": "research-agent"})) if err != nil { t.Fatalf("handleDelegate: %v", err) } @@ -169,11 +181,32 @@ func TestHandleDelegate_RequiresJWT(t *testing.T) { } } +// TestHandleDelegate_TokenScopeRejected pins the Copilot review fix: +// a token issued under a policy whose allowlist does NOT include +// aflock_delegate must be rejected before any propagation/JWT-mint +// side effect. Regression test for issue #149 review. +func TestHandleDelegate_TokenScopeRejected(t *testing.T) { + pol := policyWithSublayout() + // Drop aflock_delegate from the allowlist — the token issued from + // this policy then scopes the caller out of the new tool. + pol.Tools.Allow = []string{"Read", "Bash"} + s, tok := authedDelegateServer(t, pol) + + result, err := s.handleDelegate(context.Background(), + delegateRequest(map[string]any{"_token": tok, "sublayout_name": "research-agent"})) + if err != nil { + t.Fatalf("handleDelegate: %v", err) + } + if !result.IsError || !strings.Contains(result.Content[0].(mcp.TextContent).Text, "not permitted by token scope") { + t.Errorf("expected 'not permitted by token scope' error, got: %+v", result) + } +} + // TestHandleDelegate_MissingSublayoutName rejects empty input clearly. func TestHandleDelegate_MissingSublayoutName(t *testing.T) { s, tok := authedDelegateServer(t, policyWithSublayout()) result, err := s.handleDelegate(context.Background(), - newTestRequest(map[string]any{"_token": tok})) + delegateRequest(map[string]any{"_token": tok})) if err != nil { t.Fatalf("handleDelegate: %v", err) } @@ -186,7 +219,7 @@ func TestHandleDelegate_MissingSublayoutName(t *testing.T) { func TestHandleDelegate_UnknownSublayout(t *testing.T) { s, tok := authedDelegateServer(t, policyWithSublayout()) result, err := s.handleDelegate(context.Background(), - newTestRequest(map[string]any{"_token": tok, "sublayout_name": "nope"})) + delegateRequest(map[string]any{"_token": tok, "sublayout_name": "nope"})) if err != nil { t.Fatalf("handleDelegate: %v", err) } @@ -204,7 +237,7 @@ func TestHandleDelegate_AttenuationViolation(t *testing.T) { s, tok := authedDelegateServer(t, pol) result, err := s.handleDelegate(context.Background(), - newTestRequest(map[string]any{"_token": tok, "sublayout_name": "research-agent"})) + delegateRequest(map[string]any{"_token": tok, "sublayout_name": "research-agent"})) if err != nil { t.Fatalf("handleDelegate: %v", err) } @@ -218,7 +251,7 @@ func TestHandleDelegate_AttenuationViolation(t *testing.T) { func TestHandleDelegate_RejectsBadChildSessionID(t *testing.T) { s, tok := authedDelegateServer(t, policyWithSublayout()) result, err := s.handleDelegate(context.Background(), - newTestRequest(map[string]any{ + delegateRequest(map[string]any{ "_token": tok, "sublayout_name": "research-agent", "child_session_id": "../../etc/passwd", @@ -238,7 +271,7 @@ func TestHandleDelegate_HappyPath(t *testing.T) { s, tok := authedDelegateServer(t, policyWithSublayout()) result, err := s.handleDelegate(context.Background(), - newTestRequest(map[string]any{"_token": tok, "sublayout_name": "research-agent"})) + delegateRequest(map[string]any{"_token": tok, "sublayout_name": "research-agent"})) if err != nil { t.Fatalf("handleDelegate: %v", err) } diff --git a/internal/mcp/server.go b/internal/mcp/server.go index 04bee65..b15ffed 100644 --- a/internal/mcp/server.go +++ b/internal/mcp/server.go @@ -1605,11 +1605,19 @@ func (s *Server) handleSignAttestation(ctx context.Context, request mcp.CallTool if err != nil { return mcp.NewToolResultError(fmt.Sprintf("Authorization denied: %v", err)), nil } + // Use the actually-invoked tool name so the paper-named alias + // (aflock_attest) and the legacy name (sign_attestation) both + // authorize correctly when present in a policy allowlist + // (Copilot review on PR #149). + invoked := request.Params.Name + if invoked == "" { + invoked = "sign_attestation" + } if claims == nil && s.authActive.Load() { - return mcp.NewToolResultError("Authorization denied: sign_attestation requires a valid JWT once auth is active (issue #40)"), nil + return mcp.NewToolResultError(fmt.Sprintf("Authorization denied: %s requires a valid JWT once auth is active (issue #40)", invoked)), nil } - if claims != nil && !auth.IsToolAllowed("sign_attestation", claims.AllowedTools, claims.DeniedTools) { - return mcp.NewToolResultError("Authorization denied: tool 'sign_attestation' not permitted by token scope"), nil + if claims != nil && !auth.IsToolAllowed(invoked, claims.AllowedTools, claims.DeniedTools) { + return mcp.NewToolResultError(fmt.Sprintf("Authorization denied: tool %q not permitted by token scope", invoked)), nil } if !s.signingEnabled { @@ -1784,6 +1792,14 @@ func (s *Server) handleDelegate(ctx context.Context, request mcp.CallToolRequest if claims == nil { return mcp.NewToolResultError("Authorization denied: aflock_delegate requires a valid JWT"), nil } + // Tool-scope check: a token issued under a policy allowlist that + // does not include the invoked name must not be able to mint child + // JWTs (Copilot review on PR #149). Use the actually-invoked name + // so the same handler covers any future alias. + invoked := request.Params.Name + if !auth.IsToolAllowed(invoked, claims.AllowedTools, claims.DeniedTools) { + return mcp.NewToolResultError(fmt.Sprintf("Authorization denied: tool %q not permitted by token scope", invoked)), nil + } if s.tokenIssuer == nil { return mcp.NewToolResultError("aflock_delegate: token issuer not initialized"), nil } @@ -1839,6 +1855,14 @@ func (s *Server) handleDelegate(ctx context.Context, request mcp.CallToolRequest // Mint the JWT first. If minting fails, no propagation record is // written — failed delegations have no observable side effects. + // + // IMPORTANT (Copilot review on PR #149): the minted token is bound + // to the child session ID and the attenuated child policy digest. + // It is NOT usable against THIS server's validateJWT, which checks + // every call against the parent session/policy. The intended + // consumer is a downstream aflock instance that loads the child + // policy and runs the child session. Callers handing this token + // back to the parent socket will be rejected — by design. agentID := "" identityHash := "" if s.agentIdentity != nil { @@ -1862,6 +1886,7 @@ func (s *Server) handleDelegate(ctx context.Context, request mcp.CallToolRequest "parent_session": s.sessionID, "child_session_id": childSessionID, "child_jwt": childJWT, + "child_jwt_note": "Token is bound to child_session_id + attenuated child policy. Present it to a downstream aflock instance loaded with the child policy, not to this server.", "limits": matched.Limits, } data, _ := json.MarshalIndent(result, "", " ")