Cyrene_v2/internal/runtime/loop/tool_policy_gating_test.go

package loop

import (
	"context"
	"strings"
	"sync"
	"testing"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/capability/registry"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/capability/restrict"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/tool"
)

// newPolicyDispatcher builds a fresh registry + dispatcher and registers one
// mock tool per name, returning both so a test can set the group/stranger
// markers on the registry before exercising the design/11 §3 security gate.
func newPolicyDispatcher(t *testing.T, names []string) (*registry.Registry, *tool.Dispatcher) {
	t.Helper()
	reg := registry.New()
	d, err := tool.New(reg)
	if err != nil {
		t.Fatalf("tool.New: %v", err)
	}
	for _, name := range names {
		if err := d.Register(&mockTool{name: name}); err != nil {
			t.Fatalf("Register %q: %v", name, err)
		}
	}
	return reg, d
}

// expectToolNames asserts got carries exactly the tool names in want, any order.
// The registry List iterates a map, so registration order is not preserved;
// a set comparison keeps these assertions order-independent.
func expectToolNames(t *testing.T, got []model.ToolSpec, want ...string) {
	t.Helper()
	if len(got) != len(want) {
		t.Fatalf("tool schema count = %d, want %d (%v)", len(got), len(want), want)
	}
	wantSet := make(map[string]bool, len(want))
	for _, n := range want {
		wantSet[n] = true
	}
	for _, s := range got {
		if !wantSet[s.Name] {
			t.Fatalf("unexpected tool %q in schemas %v", s.Name, toolNames(got))
		}
	}
}

// TestToolSchemas_RestrictedEnv_FiltersByMarkers proves a group/stranger session
// is only shown the tools whose editable marker allows that environment, while
// the owner session keeps the full set (design/11 §3).
func TestToolSchemas_RestrictedEnv_FiltersByMarkers(t *testing.T) {
	t.Parallel()
	reg, d := newPolicyDispatcher(t, []string{"groupOnly", "strangerOnly", "both", "ownerOnly"})
	mustSetPolicy(t, reg, "groupOnly", registry.ToolPolicy{AllowInGroup: true})
	mustSetPolicy(t, reg, "strangerOnly", registry.ToolPolicy{AllowInStranger: true})
	mustSetPolicy(t, reg, "both", registry.ToolPolicy{AllowInGroup: true, AllowInStranger: true})
	// ownerOnly stays default-closed (zero marker).

	l := New(Options{Tools: d})

	group := l.toolSchemas(WithSessionEnv(context.Background(),
		SessionEnv{ActorEnv: restrict.ActorGroupMember, MemWritable: false}), nil)
	expectToolNames(t, group, "groupOnly", "both")

	stranger := l.toolSchemas(WithSessionEnv(context.Background(),
		SessionEnv{ActorEnv: restrict.ActorStranger, MemWritable: false}), nil)
	expectToolNames(t, stranger, "strangerOnly", "both")

	owner := l.toolSchemas(WithSessionEnv(context.Background(), OwnerEnv()), nil)
	expectToolNames(t, owner, "groupOnly", "strangerOnly", "both", "ownerOnly")
}

// TestToolSchemas_DefaultClosed_UnmarkedToolHidden proves the defaulted-closed
// marker: an unmarked tool is not advertised to a group/stranger session, so a
// session cannot even see (let alone call) an owner-only capability.
func TestToolSchemas_DefaultClosed_UnmarkedToolHidden(t *testing.T) {
	t.Parallel()
	reg, d := newPolicyDispatcher(t, []string{"secret", "shared"})
	// Only shared is opened to groups; secret is default-closed.
	mustSetPolicy(t, reg, "shared", registry.ToolPolicy{AllowInGroup: true})
	l := New(Options{Tools: d})

	ctx := WithSessionEnv(context.Background(),
		SessionEnv{ActorEnv: restrict.ActorGroupMember, MemWritable: false})
	expectToolNames(t, l.toolSchemas(ctx, nil), "shared")
}

// TestRun_Restricted_NotAllowed_BlockedAndAudited proves a group session calling
// an owner-only tool never executes it: the call is fed back as a soft
// not_allowed result (no panic/crash) and an AuditEvent is recorded for the
// rejection.
func TestRun_Restricted_NotAllowed_BlockedAndAudited(t *testing.T) {
	t.Parallel()
	call := `{"tool_calls":[{"name":"ownerOnly","args":{}}]}`
	ad := &mockAdapter{chat: []mgResp{
		{resp: model.Response{Content: call}},
		{resp: model.Response{Content: "done"}},
	}}
	reg, d := newPolicyDispatcher(t, []string{"ownerOnly", "shared"})
	mustSetPolicy(t, reg, "shared", registry.ToolPolicy{AllowInGroup: true})

	a := newTestAgent(t, ad, d)
	var mu sync.Mutex
	var audited []AuditEvent
	l := New(Options{Agent: a, MaxRounds: 8, OnAudit: func(ev AuditEvent) {
		mu.Lock()
		audited = append(audited, ev)
		mu.Unlock()
	}})

	ctx := WithSessionEnv(context.Background(),
		SessionEnv{ActorEnv: restrict.ActorGroupMember, MemWritable: false})
	res, err := l.Run(ctx, "sess", []model.Message{model.TextMessage(model.RoleUser, "go")})
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if res.Text != "done" {
		t.Fatalf("Text = %q, want done", res.Text)
	}
	if res.ToolCallsExecuted != 0 {
		t.Fatalf("ToolCallsExecuted = %d, want 0 (ownerOnly blocked in group)", res.ToolCallsExecuted)
	}

	last := ad.reqs[len(ad.reqs)-1].Messages
	lastMsg := last[len(last)-1]
	if !strings.Contains(lastMsg.Content[0].Text, "not_allowed") ||
		!strings.Contains(lastMsg.Content[0].Text, "ownerOnly") {
		t.Fatalf("expected not_allowed soft result, got %q", lastMsg.Content[0].Text)
	}

	mu.Lock()
	defer mu.Unlock()
	if len(audited) != 1 {
		t.Fatalf("audit events = %d, want 1", len(audited))
	}
	if audited[0].What != auditToolRejectNotAllowed {
		t.Fatalf("audit what = %q, want %q", audited[0].What, auditToolRejectNotAllowed)
	}
	if audited[0].Who != string(restrict.ActorGroupMember) {
		t.Fatalf("audit who = %q, want %q", audited[0].Who, restrict.ActorGroupMember)
	}
	if !strings.Contains(audited[0].Result, "ownerOnly") {
		t.Fatalf("audit result = %q, want to name ownerOnly", audited[0].Result)
	}
}

// TestRun_Restricted_AllowedTool_Executes proves a group session may invoke a
// tool whose AllowInGroup marker is set, and that it runs normally.
func TestRun_Restricted_AllowedTool_Executes(t *testing.T) {
	t.Parallel()
	call := `{"tool_calls":[{"name":"shared","args":{"q":"hi"}}]}`
	ad := &mockAdapter{chat: []mgResp{
		{resp: model.Response{Content: call}},
		{resp: model.Response{Content: "done"}},
	}}
	reg, d := newPolicyDispatcher(t, []string{"shared", "ownerOnly"})
	mustSetPolicy(t, reg, "shared", registry.ToolPolicy{AllowInGroup: true})

	a := newTestAgent(t, ad, d)
	l := New(Options{Agent: a, MaxRounds: 8})

	ctx := WithSessionEnv(context.Background(),
		SessionEnv{ActorEnv: restrict.ActorGroupMember, MemWritable: false})
	res, err := l.Run(ctx, "sess", []model.Message{model.TextMessage(model.RoleUser, "go")})
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if res.ToolCallsExecuted != 1 {
		t.Fatalf("ToolCallsExecuted = %d, want 1 (shared allowed in group)", res.ToolCallsExecuted)
	}
	if res.Text != "done" {
		t.Fatalf("Text = %q, want done", res.Text)
	}
}

// TestRun_Owner_AllToolsExecute proves the owner environment is unaffected by the
// default-closed marker: the owner may invoke an owner-only (unmarked) tool.
func TestRun_Owner_AllToolsExecute(t *testing.T) {
	t.Parallel()
	call := `{"tool_calls":[{"name":"ownerOnly","args":{}}]}`
	ad := &mockAdapter{chat: []mgResp{
		{resp: model.Response{Content: call}},
		{resp: model.Response{Content: "done"}},
	}}
	_, d := newPolicyDispatcher(t, []string{"ownerOnly"})
	// Default-closed marker (never opened): owner must still be able to call it.
	a := newTestAgent(t, ad, d)
	l := New(Options{Agent: a, MaxRounds: 8})

	res, err := l.Run(WithSessionEnv(context.Background(), OwnerEnv()), "sess",
		[]model.Message{model.TextMessage(model.RoleUser, "go")})
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if res.ToolCallsExecuted != 1 {
		t.Fatalf("ToolCallsExecuted = %d, want 1 (owner unaffected by default-closed)", res.ToolCallsExecuted)
	}
	if res.Text != "done" {
		t.Fatalf("Text = %q, want done", res.Text)
	}
}

// TestRun_Stranger_NotAllowed_Blocked proves the gate also applies to the
// stranger private-chat environment, and that a group-only tool is denied there.
func TestRun_Stranger_NotAllowed_Blocked(t *testing.T) {
	t.Parallel()
	call := `{"tool_calls":[{"name":"groupOnly","args":{}}]}`
	ad := &mockAdapter{chat: []mgResp{
		{resp: model.Response{Content: call}},
		{resp: model.Response{Content: "done"}},
	}}
	reg, d := newPolicyDispatcher(t, []string{"groupOnly", "strangerOk"})
	mustSetPolicy(t, reg, "groupOnly", registry.ToolPolicy{AllowInGroup: true})
	mustSetPolicy(t, reg, "strangerOk", registry.ToolPolicy{AllowInStranger: true})

	a := newTestAgent(t, ad, d)
	l := New(Options{Agent: a, MaxRounds: 8})

	ctx := WithSessionEnv(context.Background(),
		SessionEnv{ActorEnv: restrict.ActorStranger, MemWritable: false})
	res, err := l.Run(ctx, "sess", []model.Message{model.TextMessage(model.RoleUser, "go")})
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if res.ToolCallsExecuted != 0 {
		t.Fatalf("ToolCallsExecuted = %d, want 0 (groupOnly denied to stranger)", res.ToolCallsExecuted)
	}
	last := ad.reqs[len(ad.reqs)-1].Messages
	if !strings.Contains(last[len(last)-1].Content[0].Text, "not_allowed") {
		t.Fatalf("expected not_allowed soft result, got %q", last[len(last)-1].Content[0].Text)
	}
}

// mustSetPolicy applies a ToolPolicy to the named tool, failing the test on error.
func mustSetPolicy(t *testing.T, reg *registry.Registry, name string, pol registry.ToolPolicy) {
	t.Helper()
	if err := reg.SetPolicy(name, pol); err != nil {
		t.Fatalf("SetPolicy %q: %v", name, err)
	}
}