mirror of
https://github.com/RuoJi6/ARTEX.git
synced 2026-10-28 17:13:50 +08:00
feat: filter approval records by status and decision source
This commit is contained in:
ji ruo committed
1 parent
80a085d157
commit
a91c0a7223
8 files changed
+431
-47
No files matched your search
+32
-13
@@ -225,8 +225,13 @@ func scanInterceptApprovalRow(rows interface{ Scan(...any) error }, r *Intercept
|
||||
)
|
||||
}
|
||||
|
||||
// Keep legacy rows without decision_source consistent with their displayed source.
|
||||
const approvalDecisionSource = `COALESCE(NULLIF(ip.decision_source,''), CASE
|
||||
WHEN ip.rule_id IS NOT NULL THEN 'rule'
|
||||
WHEN ip.reason LIKE '[模型]%' THEN 'model' ELSE 'unknown' END)`
|
||||
|
||||
const approvalRowColumns = `ip.id, ip.rule_id, ip.conversation_id, ip.task_id, ip.agent_name,
|
||||
ip.tool_name, ip.tool_input, ip.status, ip.reason, ip.decided_at, ip.created_at, ip.decision_source,
|
||||
ip.tool_name, ip.tool_input, ip.status, ip.reason, ip.decided_at, ip.created_at, ` + approvalDecisionSource + `,
|
||||
COALESCE(c.title,'') AS conv_title,
|
||||
COALESCE(c.agent_key,'') AS conv_agent_key,
|
||||
COALESCE(ir.name,'') AS rule_name`
|
||||
@@ -279,11 +284,16 @@ func (d *DB) ListAllIntercepts(limit int) ([]InterceptApprovalRow, error) {
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// ListAllInterceptsPage returns one 1-based page of approval records and the
|
||||
// total number of records. Records are ordered newest first, matching the
|
||||
// legacy capped list returned by ListAllIntercepts.
|
||||
func (d *DB) ListAllInterceptsPage(page, size int) ([]InterceptApprovalRow, int, error) {
|
||||
return d.listInterceptsPage("", page, size)
|
||||
// InterceptApprovalFilter combines exact status and decision-source filters.
|
||||
// Empty fields include all values.
|
||||
type InterceptApprovalFilter struct {
|
||||
Status string
|
||||
DecisionSource string
|
||||
}
|
||||
|
||||
// ListAllInterceptsPage returns one 1-based page and the total matching count.
|
||||
func (d *DB) ListAllInterceptsPage(page, size int, filter InterceptApprovalFilter) ([]InterceptApprovalRow, int, error) {
|
||||
return d.listInterceptsPage("", page, size, filter)
|
||||
}
|
||||
|
||||
// ListTaskIntercepts returns all intercept_pending rows for a specific task (newest first).
|
||||
@@ -305,11 +315,11 @@ func (d *DB) ListTaskIntercepts(taskID string) ([]InterceptApprovalRow, error) {
|
||||
}
|
||||
|
||||
// ListTaskInterceptsPage is the paginated variant of ListTaskIntercepts.
|
||||
func (d *DB) ListTaskInterceptsPage(taskID string, page, size int) ([]InterceptApprovalRow, int, error) {
|
||||
return d.listInterceptsPage(taskID, page, size)
|
||||
func (d *DB) ListTaskInterceptsPage(taskID string, page, size int, filter InterceptApprovalFilter) ([]InterceptApprovalRow, int, error) {
|
||||
return d.listInterceptsPage(taskID, page, size, filter)
|
||||
}
|
||||
|
||||
func (d *DB) listInterceptsPage(taskID string, page, size int) ([]InterceptApprovalRow, int, error) {
|
||||
func (d *DB) listInterceptsPage(taskID string, page, size int, filter InterceptApprovalFilter) ([]InterceptApprovalRow, int, error) {
|
||||
if page < 1 {
|
||||
page = 1
|
||||
}
|
||||
@@ -321,11 +331,20 @@ func (d *DB) listInterceptsPage(taskID string, page, size int) ([]InterceptAppro
|
||||
}
|
||||
offset := (page - 1) * size
|
||||
|
||||
where := ""
|
||||
conditions := []string{}
|
||||
args := []any{}
|
||||
if taskID != "" {
|
||||
where = " WHERE ip.task_id=$1"
|
||||
args = append(args, taskID)
|
||||
add := func(column, value string) {
|
||||
if value != "" {
|
||||
args = append(args, value)
|
||||
conditions = append(conditions, column+"=$"+fmt.Sprint(len(args)))
|
||||
}
|
||||
}
|
||||
add("ip.task_id", taskID)
|
||||
add("ip.status", filter.Status)
|
||||
add(approvalDecisionSource, filter.DecisionSource)
|
||||
where := ""
|
||||
if len(conditions) > 0 {
|
||||
where = " WHERE " + strings.Join(conditions, " AND ")
|
||||
}
|
||||
var total int
|
||||
if err := d.QueryRow("SELECT COUNT(*) FROM intercept_pending ip"+where, args...).Scan(&total); err != nil {
|
||||
|
||||
@@ -0,0 +1,141 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"slices"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestInterceptApprovalFilters(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.Close() })
|
||||
scope := "approval-filter-" + t.Name()
|
||||
t.Cleanup(func() { _, _ = d.Exec(`DELETE FROM intercept_pending WHERE task_id IN ($1,$2)`, scope, scope+"-other") })
|
||||
rule, err := d.CreateInterceptRule("filter fixture", "tool_name", "string", "Bash", "deny", "fixture", 1, true, false, 0, "deny")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _, _ = d.Exec(`DELETE FROM intercept_rules WHERE id=$1`, rule.ID) })
|
||||
filter := InterceptApprovalFilter{Status: "denied", DecisionSource: "model"}
|
||||
_, baseline, err := d.ListAllInterceptsPage(1, 20, filter)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
expected := []int64{}
|
||||
for _, task := range []string{scope, scope + "-other"} {
|
||||
copies := 3
|
||||
if task != scope {
|
||||
copies = 1
|
||||
}
|
||||
for _, status := range []string{"pending", "allowed", "denied", "timeout"} {
|
||||
for _, source := range []string{"model", "rule", "unknown"} {
|
||||
for i := 0; i < copies; i++ {
|
||||
var ruleID int64
|
||||
reason := "legacy reason"
|
||||
if source == "rule" {
|
||||
ruleID = rule.ID
|
||||
}
|
||||
if source == "model" {
|
||||
reason = "[模型] fixture"
|
||||
}
|
||||
id, err := d.CreateDecidedIntercept(ruleID, 0, task, "test", "Bash", []byte(`{"command":"fixture"}`), status, reason)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// Legacy source fields and tied timestamps must filter/page consistently.
|
||||
if i == 0 {
|
||||
if _, err = d.Exec(`UPDATE intercept_pending SET decision_source='' WHERE id=$1`, id); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if task == scope && status == "denied" && source == "model" {
|
||||
expected = append(expected, id)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if _, err := d.Exec(`UPDATE intercept_pending SET created_at='2099-01-01' WHERE task_id IN ($1,$2)`, scope, scope+"-other"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cases := []struct {
|
||||
name string
|
||||
filter InterceptApprovalFilter
|
||||
want int
|
||||
}{
|
||||
{"all", InterceptApprovalFilter{}, 36},
|
||||
}
|
||||
for _, s := range []string{"pending", "allowed", "denied", "timeout"} {
|
||||
cases = append(cases, struct {
|
||||
name string
|
||||
filter InterceptApprovalFilter
|
||||
want int
|
||||
}{s, InterceptApprovalFilter{Status: s}, 9})
|
||||
}
|
||||
for _, s := range []string{"model", "rule", "unknown"} {
|
||||
cases = append(cases, struct {
|
||||
name string
|
||||
filter InterceptApprovalFilter
|
||||
want int
|
||||
}{s, InterceptApprovalFilter{DecisionSource: s}, 12})
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
rows, total, err := d.ListTaskInterceptsPage(scope, 1, 100, tc.filter)
|
||||
if err != nil || total != tc.want || len(rows) != tc.want {
|
||||
t.Fatalf("count=%d rows=%d err=%v", total, len(rows), err)
|
||||
}
|
||||
for _, r := range rows {
|
||||
if r.TaskID == nil || *r.TaskID != scope || tc.filter.Status != "" && r.Status != tc.filter.Status || tc.filter.DecisionSource != "" && r.DecisionSource != tc.filter.DecisionSource {
|
||||
t.Fatalf("nonmatching row: %+v", r)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
slices.Reverse(expected)
|
||||
var seen []int64
|
||||
for page := 1; page <= 2; page++ {
|
||||
rows, total, err := d.ListTaskInterceptsPage(scope, page, 2, filter)
|
||||
if err != nil || total != 3 {
|
||||
t.Fatalf("page %d total=%d err=%v", page, total, err)
|
||||
}
|
||||
for _, r := range rows {
|
||||
seen = append(seen, r.ID)
|
||||
}
|
||||
}
|
||||
if !slices.Equal(seen, expected) {
|
||||
t.Fatalf("unstable combined-filter pagination: got %v want %v", seen, expected)
|
||||
}
|
||||
rows, total, err := d.ListTaskInterceptsPage(scope, 3, 2, filter)
|
||||
if err != nil || len(rows) != 0 || total != 3 {
|
||||
t.Fatalf("empty page count=%d rows=%d err=%v", total, len(rows), err)
|
||||
}
|
||||
_, total, err = d.ListAllInterceptsPage(1, 100, filter)
|
||||
if err != nil || total != baseline+4 {
|
||||
t.Fatalf("global filter count=%d want=%d err=%v", total, baseline+4, err)
|
||||
}
|
||||
// Values must remain SQL parameters even for direct store callers.
|
||||
rows, total, err = d.ListTaskInterceptsPage(scope, 1, 20, InterceptApprovalFilter{Status: "denied' OR 1=1 --"})
|
||||
if err != nil || total != 0 || len(rows) != 0 {
|
||||
t.Fatalf("invalid value escaped filter: %d %v", total, err)
|
||||
}
|
||||
// A rule's persisted source survives deletion of its associated rule.
|
||||
_, err = d.Exec(`UPDATE intercept_pending SET decision_source='rule' WHERE task_id=$1 AND rule_id=$2`, scope, rule.ID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = d.DeleteInterceptRule(rule.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
rows, total, err = d.ListTaskInterceptsPage(scope, 1, 100, InterceptApprovalFilter{DecisionSource: "rule"})
|
||||
if err != nil || total != 12 {
|
||||
t.Fatalf("deleted-rule source count=%d err=%v", total, err)
|
||||
}
|
||||
for _, r := range rows {
|
||||
if r.RuleID != nil || r.DecisionSource != "rule" {
|
||||
t.Fatal("deleted rule lost its source")
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user