Merge pull request #178 from BaaRai/feat/target-gate-schema-request

feat(intercept): target-gate schema, opt-in flag and approval requests (S3b 1.2-1.4)
This commit is contained in:
BaaRai
2026-11-10 23:35:55 +01:00
committed by GitHub
24 changed files with 3496 additions and 61 deletions
+1 -1
View File
@@ -70,7 +70,7 @@ func TestCaptureApprovalLifecycle(t *testing.T) {
t.Error(err)
return
}
if err := ic.Decide(detail.ID, tc.approve); err != nil {
if err := ic.Decide(detail.ID, tc.approve, "", ""); err != nil {
t.Error(err)
}
})
+51 -15
View File
@@ -58,20 +58,24 @@ type Conversation struct {
PinnedAt *time.Time `json:"pinned_at,omitempty"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
// TargetGate is the per-conversation opt-in of the approval gate for asset
// targets (a chat has no task). false (default) = no gate.
TargetGate bool `json:"target_gate"`
}
// ConversationPatch updates only the fields whose pointers are non-nil.
type ConversationPatch struct {
Title *string
Pinned *bool
Title *string
Pinned *bool
TargetGate *bool
}
const convCols = `id, agent_key, title, llm_profile_id, pinned_at, created_at, updated_at`
const convCols = `id, agent_key, title, llm_profile_id, pinned_at, created_at, updated_at, target_gate`
func scanConv(row interface{ Scan(...any) error }) (Conversation, error) {
var c Conversation
var pinnedAt sql.NullTime
err := row.Scan(&c.ID, &c.AgentKey, &c.Title, &c.LLMProfileID, &pinnedAt, &c.CreatedAt, &c.UpdatedAt)
err := row.Scan(&c.ID, &c.AgentKey, &c.Title, &c.LLMProfileID, &pinnedAt, &c.CreatedAt, &c.UpdatedAt, &c.TargetGate)
if pinnedAt.Valid {
c.Pinned = true
c.PinnedAt = &pinnedAt.Time
@@ -82,6 +86,12 @@ func scanConv(row interface{ Scan(...any) error }) (Conversation, error) {
// CreateConversation opens a new chat thread for agentKey with an initial title.
// llmProfileID may be nil to use the globally active profile.
func (d *DB) CreateConversation(agentKey, title string, llmProfileID *int64) (*Conversation, error) {
return d.CreateConversationGated(agentKey, title, llmProfileID, false)
}
// CreateConversationGated is CreateConversation with the target-gate opt-in set
// at creation.
func (d *DB) CreateConversationGated(agentKey, title string, llmProfileID *int64, targetGate bool) (*Conversation, error) {
tx, err := d.Begin()
if err != nil {
return nil, err
@@ -92,8 +102,8 @@ func (d *DB) CreateConversation(agentKey, title string, llmProfileID *int64) (*C
// by this transaction before it takes a profile lock, matching DeleteProfile's
// child-row -> profile-row protocol.
c, err := scanConv(tx.QueryRow(`
INSERT INTO conversations(agent_key, title, llm_profile_id) VALUES ($1, $2, NULL)
RETURNING `+convCols, agentKey, title))
INSERT INTO conversations(agent_key, title, llm_profile_id, target_gate) VALUES ($1, $2, NULL, $3)
RETURNING `+convCols, agentKey, title, targetGate))
if err != nil {
return nil, err
}
@@ -169,7 +179,7 @@ func (d *DB) GetConversation(id int64) (*Conversation, error) {
return &c, nil
}
// UpdateConversation applies a partial title/pin mutation and returns the updated
// UpdateConversation applies a partial title/pin/target-gate mutation and returns the updated
// row. Pinning an already-pinned conversation preserves its original pin order.
func (d *DB) UpdateConversation(id int64, patch ConversationPatch) (*Conversation, error) {
c, err := scanConv(d.QueryRow(`UPDATE conversations SET
@@ -178,9 +188,10 @@ func (d *DB) UpdateConversation(id int64, patch ConversationPatch) (*Conversatio
WHEN $4::boolean IS NULL THEN pinned_at
WHEN $4::boolean THEN COALESCE(pinned_at, now())
ELSE NULL
END
END,
target_gate = CASE WHEN $5::boolean THEN $6 ELSE target_gate END
WHERE id=$1
RETURNING `+convCols, id, patch.Title != nil, patch.Title, patch.Pinned))
RETURNING `+convCols, id, patch.Title != nil, patch.Title, patch.Pinned, patch.TargetGate != nil, patch.TargetGate))
if err == sql.ErrNoRows {
return nil, nil
}
@@ -203,19 +214,37 @@ func (d *DB) TouchConversation(id int64) error {
return err
}
// DeleteConversation removes a thread; its activities cascade via FK.
// DeleteConversation removes a thread; its activities cascade via FK. Its
// target pins are not rows of the thread: they are revoked in the same
// transaction and stay as history.
func (d *DB) DeleteConversation(id int64) error {
_, err := d.Exec(`DELETE FROM conversations WHERE id=$1`, id)
return err
tx, err := d.Begin()
if err != nil {
return err
}
defer tx.Rollback()
if err := revokeConversationPins(tx, []int64{id}); err != nil {
return err
}
if _, err := tx.Exec(`DELETE FROM conversations WHERE id=$1`, id); err != nil {
return err
}
return tx.Commit()
}
// DeleteConversations removes existing threads in one statement and returns the
// ids that were actually present. Child activities and trigger runs cascade.
// ids that were actually present. Child activities and trigger runs cascade. The
// target pins of the removed threads are revoked in the same transaction.
func (d *DB) DeleteConversations(ids []int64) ([]int64, error) {
if len(ids) == 0 {
return []int64{}, nil
}
rows, err := d.Query(`DELETE FROM conversations WHERE id=ANY($1::bigint[]) RETURNING id`, ids)
tx, err := d.Begin()
if err != nil {
return nil, err
}
defer tx.Rollback()
rows, err := tx.Query(`DELETE FROM conversations WHERE id=ANY($1::bigint[]) RETURNING id`, ids)
if err != nil {
return nil, err
}
@@ -228,7 +257,14 @@ func (d *DB) DeleteConversations(ids []int64) ([]int64, error) {
}
deleted = append(deleted, id)
}
return deleted, rows.Err()
if err := rows.Err(); err != nil {
return nil, err
}
rows.Close()
if err := revokeConversationPins(tx, deleted); err != nil {
return nil, err
}
return deleted, tx.Commit()
}
// AppendConvActivity records one step of a conversation (human message or an agent
+5 -2
View File
@@ -114,8 +114,11 @@ func (d *DB) CreateFindingRetest(ctx context.Context, findingID int64, notes str
if runes := []rune(title); len(runes) > 100 {
title = string(runes[:100])
}
c, err := scanConv(tx.QueryRowContext(ctx, `INSERT INTO conversations(agent_key,title) VALUES ($1,$2) RETURNING `+convCols,
FindingRetestAgentKey, fmt.Sprintf("Retest #%d · %s", findingID, title)))
// The retest conversation takes the target-gate opt-in of the finding's task.
// Pins are never copied: a pin belongs to the scope it was approved in.
c, err := scanConv(tx.QueryRowContext(ctx, `INSERT INTO conversations(agent_key,title,target_gate) VALUES ($1,$2,
COALESCE((SELECT t.target_gate FROM findings f JOIN tasks t ON t.id=f.task_id WHERE f.id=$3), false)) RETURNING `+convCols,
FindingRetestAgentKey, fmt.Sprintf("Retest #%d · %s", findingID, title), findingID))
if err != nil {
return nil, nil, false, err
}
+46 -2
View File
@@ -43,6 +43,35 @@ type InterceptAudit struct {
Output string `json:"output,omitempty"`
OutputTruncated bool `json:"output_truncated,omitempty"`
ExecutionEndedAt *time.Time `json:"execution_ended_at,omitempty"`
// Targets is set on target approval requests only (target_key IS NOT NULL).
Targets []InterceptTarget `json:"targets,omitempty"`
// MergedAsk marks a target request that also carries the ask of a rule or of
// the judge: one row, one decision for both.
MergedAsk bool `json:"merged_ask,omitempty"`
// Approver and ApproverAddr identify the human behind a decision: the subject
// of the session token and the remote address of the request. decided_by keeps
// the actor class (human).
Approver string `json:"approver,omitempty"`
ApproverAddr string `json:"approver_addr,omitempty"`
}
// InterceptTarget is one asset target of a target approval request: what was
// asked for, the addresses shown to the human, and the pin it relates to.
type InterceptTarget struct {
Host string `json:"host"`
Port int `json:"port"` // 0 = port not declared
Kind string `json:"kind,omitempty"`
ResolvedIPs []string `json:"resolved_ips"`
Resolver string `json:"resolver,omitempty"`
ResolvedAt time.Time `json:"resolved_at"`
PinID int64 `json:"pin_id,omitempty"`
PinVersion int `json:"pin_version,omitempty"`
PinnedIPs []string `json:"pinned_ips,omitempty"`
Reason string `json:"reason,omitempty"` // no_pin | new_ip | revoked
Outcome string `json:"outcome"` // requested | pinned | denied
// NameOnlyEffective: requested while the upstream egress proxy was set; a pin
// created from this request covers the name only.
NameOnlyEffective bool `json:"name_only_effective,omitempty"`
}
type InterceptDetail struct {
@@ -77,6 +106,14 @@ func (d *DB) GetInterceptDetail(id int64) (*InterceptDetail, error) {
// decidedBy records the settling actor (human | rule | model | system); it is
// supplied by the caller because status alone cannot tell them apart.
func (d *DB) ResolveIntercept(id int64, status, action, reason, decidedBy string) (bool, error) {
return d.ResolveInterceptBy(id, status, action, reason, decidedBy, "", "")
}
// ResolveInterceptBy is ResolveIntercept for a decision whose author is known:
// approver (subject of the session token) and remoteAddr (address of the
// request) are recorded in the row's audit when non-empty. A row without an
// audit keeps none.
func (d *DB) ResolveInterceptBy(id int64, status, action, reason, decidedBy, approver, remoteAddr string) (bool, error) {
switch decidedBy {
case "human", "rule", "model", "system":
default:
@@ -86,9 +123,16 @@ func (d *DB) ResolveIntercept(id int64, status, action, reason, decidedBy string
if action == "allow" {
execution = "awaiting_result"
}
patch, err := json.Marshal(map[string]any{
fields := map[string]any{
"effective_action": action, "decision_reason": reason, "execution_status": execution,
})
}
if approver != "" {
fields["approver"] = approver
}
if remoteAddr != "" {
fields["approver_addr"] = remoteAddr
}
patch, err := json.Marshal(fields)
if err != nil {
return false, err
}
+53
View File
@@ -558,6 +558,8 @@ CREATE TRIGGER trg_tasks_upd BEFORE UPDATE ON tasks
-- planner heartbeat trigger interval (seconds); backfills old databases. Default 300s (5min). See docs/planner-trigger-impl-plan.md
ALTER TABLE tasks ADD COLUMN IF NOT EXISTS plan_heartbeat_seconds INTEGER NOT NULL DEFAULT 300;
ALTER TABLE tasks ADD COLUMN IF NOT EXISTS plan_quiet_mode BOOLEAN NOT NULL DEFAULT true;
-- Approval gate for asset targets, opt-in per task; backfills old databases. false (default) = no gate, behaviour unchanged.
ALTER TABLE tasks ADD COLUMN IF NOT EXISTS target_gate BOOLEAN NOT NULL DEFAULT false;
ALTER TABLE tasks ADD COLUMN IF NOT EXISTS llm_rotate BOOLEAN NOT NULL DEFAULT false;
-- Concurrency-limit hold state; backfills old databases. true=queued due to the concurrency limit, auto-starts when a slot frees up.
ALTER TABLE tasks ADD COLUMN IF NOT EXISTS queued BOOLEAN NOT NULL DEFAULT false;
@@ -995,6 +997,8 @@ CREATE TABLE IF NOT EXISTS conversations (
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
ALTER TABLE conversations ADD COLUMN IF NOT EXISTS pinned_at TIMESTAMPTZ;
-- Approval gate for asset targets, opt-in per conversation (a chat has no task). false (default) = no gate.
ALTER TABLE conversations ADD COLUMN IF NOT EXISTS target_gate BOOLEAN NOT NULL DEFAULT false;
CREATE INDEX IF NOT EXISTS idx_conversations_llm_profile ON conversations(llm_profile_id) WHERE llm_profile_id IS NOT NULL;
CREATE INDEX IF NOT EXISTS idx_conversations_pinned ON conversations(pinned_at DESC) WHERE pinned_at IS NOT NULL;
DROP TRIGGER IF EXISTS trg_conversations_upd ON conversations;
@@ -1119,6 +1123,55 @@ UPDATE intercept_pending SET decision_source='model'
-- existed) | human | rule | model | system. No backfill: readers fall back to
-- decision_source for legacy rows instead of rewriting history.
ALTER TABLE intercept_pending ADD COLUMN IF NOT EXISTS decided_by TEXT NOT NULL DEFAULT '';
-- Target approval requests. target_key = '<host>:<port>' of the asset target the row asks for;
-- NULL = classic ask (and every row that predates the column). At most one open request per
-- scope (task or conversation) and target: concurrent calls to one host share one row.
ALTER TABLE intercept_pending ADD COLUMN IF NOT EXISTS target_key TEXT;
CREATE UNIQUE INDEX IF NOT EXISTS uq_intercept_pending_target
ON intercept_pending (COALESCE(task_id,''), COALESCE(conversation_id,0), target_key)
WHERE status = 'pending' AND target_key IS NOT NULL;
-- Target pins: host + port approved by a human for one task or one conversation, with the
-- addresses shown at approval time. Written only by the decide endpoint. No foreign keys on
-- task_id, conversation_id and intercept_id: a pin outlives its task, its conversation and its
-- request row as history (revoked, never deleted). status and revoked_reason carry no CHECK
-- (validated in Go, same convention as findings.status).
CREATE TABLE IF NOT EXISTS target_pins (
id BIGSERIAL PRIMARY KEY,
task_id BIGINT,
conversation_id BIGINT,
host TEXT NOT NULL, -- normalised name, IP literal or CIDR
port INTEGER NOT NULL DEFAULT 0, -- 0 = port not declared: covers every port of the host
version INTEGER NOT NULL DEFAULT 1,
pinned_ips TEXT[] NOT NULL DEFAULT '{}', -- canonical, sorted
resolver TEXT NOT NULL DEFAULT '',
resolved_at TIMESTAMPTZ NOT NULL DEFAULT now(),
status TEXT NOT NULL DEFAULT 'active', -- active | superseded | revoked
intercept_id BIGINT, -- the request row that carried the click
decided_by TEXT NOT NULL DEFAULT '', -- approver: subject of the session token
decided_from TEXT NOT NULL DEFAULT '', -- remote address of the click
-- Approved while the upstream egress proxy was set: the proxy chooses the address, so the
-- pin covers the name only and its addresses are not compared. Stays true if the proxy is
-- removed later; enforcing addresses again takes a new approval.
name_only_effective BOOLEAN NOT NULL DEFAULT false,
revoked_at TIMESTAMPTZ,
revoked_reason TEXT NOT NULL DEFAULT '', -- operator | task_deleted | conversation_deleted
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
CREATE UNIQUE INDEX IF NOT EXISTS uq_target_pins_active
ON target_pins (COALESCE(task_id,0), COALESCE(conversation_id,0), host, port)
WHERE status = 'active';
CREATE INDEX IF NOT EXISTS idx_target_pins_scope ON target_pins (task_id, conversation_id, status);
-- Append-only history of the target-gate exemption list (settings key target_gate_exempt_hosts).
CREATE TABLE IF NOT EXISTS target_gate_exemption_log (
id BIGSERIAL PRIMARY KEY,
changed_at TIMESTAMPTZ NOT NULL DEFAULT now(),
changed_by TEXT NOT NULL DEFAULT '',
before JSONB,
after JSONB,
digest TEXT NOT NULL DEFAULT ''
);
-- =====================================================================
-- L. Finding persistence
+393
View File
@@ -0,0 +1,393 @@
package db
import (
"database/sql"
"encoding/json"
"errors"
"fmt"
"sort"
"strconv"
"time"
)
// TargetScopeRef names what a target pin or a target approval request is bound
// to: one task or one conversation, never both.
type TargetScopeRef struct {
TaskID int64
ConversationID int64
}
func (s TargetScopeRef) check() error {
if (s.TaskID > 0) == (s.ConversationID > 0) {
return fmt.Errorf("target scope must name one task or one conversation, got task %d and conversation %d", s.TaskID, s.ConversationID)
}
return nil
}
// pendingScope is the scope as intercept_pending stores it: task_id is TEXT
// there, and the column of the other kind is NULL.
func (s TargetScopeRef) pendingScope() (taskID *string, convID *int64) {
if s.TaskID > 0 {
v := strconv.FormatInt(s.TaskID, 10)
return &v, nil
}
return nil, &s.ConversationID
}
// TargetKey is the key of an asset target inside its scope: host and port. The
// port is always the part after the last colon, so the key stays unambiguous
// for IPv6 literals and CIDRs.
func TargetKey(host string, port int) string {
return host + ":" + strconv.Itoa(port)
}
// Target pin states and revocation reasons. The columns carry no CHECK; these
// are the only values written.
const (
TargetPinActive = "active"
TargetPinSuperseded = "superseded"
TargetPinRevoked = "revoked"
TargetPinRevokedByOperator = "operator"
TargetPinRevokedTaskDeleted = "task_deleted"
TargetPinRevokedConversationDelete = "conversation_deleted"
)
// TargetPin is one row of target_pins: a host and port a human approved for one
// scope, with the addresses that were shown at approval time. Port 0 means the
// port was not declared and covers every port of the host.
type TargetPin struct {
ID int64 `json:"id"`
TaskID *int64 `json:"task_id"`
ConversationID *int64 `json:"conversation_id"`
Host string `json:"host"`
Port int `json:"port"`
Version int `json:"version"`
PinnedIPs []string `json:"pinned_ips"`
Resolver string `json:"resolver"`
ResolvedAt time.Time `json:"resolved_at"`
Status string `json:"status"`
InterceptID *int64 `json:"intercept_id"`
DecidedBy string `json:"decided_by"`
DecidedFrom string `json:"decided_from"`
// NameOnlyEffective: approved while the upstream egress proxy was set, so the
// pin covers the name and its addresses are not compared. It stays set when
// the proxy is removed later.
NameOnlyEffective bool `json:"name_only_effective"`
RevokedAt *time.Time `json:"revoked_at,omitempty"`
RevokedReason string `json:"revoked_reason,omitempty"`
CreatedAt time.Time `json:"created_at"`
}
const targetPinCols = `id, task_id, conversation_id, host, port, version, array_to_json(pinned_ips)::text, resolver, resolved_at, status, intercept_id, decided_by, decided_from, name_only_effective, revoked_at, revoked_reason, created_at`
// targetPinScope matches the expressions of uq_target_pins_active.
const targetPinScope = `COALESCE(task_id,0)=$1 AND COALESCE(conversation_id,0)=$2`
func scanTargetPin(row interface{ Scan(...any) error }) (*TargetPin, error) {
var p TargetPin
var ips string
if err := row.Scan(&p.ID, &p.TaskID, &p.ConversationID, &p.Host, &p.Port, &p.Version, &ips, &p.Resolver, &p.ResolvedAt,
&p.Status, &p.InterceptID, &p.DecidedBy, &p.DecidedFrom, &p.NameOnlyEffective, &p.RevokedAt, &p.RevokedReason, &p.CreatedAt); err != nil {
return nil, err
}
if err := json.Unmarshal([]byte(ips), &p.PinnedIPs); err != nil {
return nil, err
}
return &p, nil
}
func queryTargetPins(q interface {
Query(string, ...any) (*sql.Rows, error)
}, query string, args ...any) ([]*TargetPin, error) {
rows, err := q.Query(query, args...)
if err != nil {
return nil, err
}
defer rows.Close()
out := []*TargetPin{}
for rows.Next() {
p, err := scanTargetPin(rows)
if err != nil {
return nil, err
}
out = append(out, p)
}
return out, rows.Err()
}
// ActiveTargetPins returns the active pins that may cover the target
// (host, port) in scope, the pin on the exact port first. A target on a declared
// port has two candidates: the pin on that port and the pin on port 0. A target
// whose port is not declared (0) has one: a pin on a specific port never covers
// it. The address sets of two pins are never merged by the caller.
func (d *DB) ActiveTargetPins(scope TargetScopeRef, host string, port int) ([]*TargetPin, error) {
if err := scope.check(); err != nil {
return nil, err
}
return queryTargetPins(d, `SELECT `+targetPinCols+` FROM target_pins
WHERE `+targetPinScope+` AND host=$3 AND port IN ($4, 0) AND status='active'
ORDER BY port DESC`, scope.TaskID, scope.ConversationID, host, port)
}
// ListTargetPins returns every pin of scope in any state, newest first.
func (d *DB) ListTargetPins(scope TargetScopeRef) ([]*TargetPin, error) {
if err := scope.check(); err != nil {
return nil, err
}
return queryTargetPins(d, `SELECT `+targetPinCols+` FROM target_pins
WHERE `+targetPinScope+` ORDER BY id DESC`, scope.TaskID, scope.ConversationID)
}
// RevokeTargetPin revokes one active pin of scope on behalf of the operator and
// returns the pin as it is afterwards (nil when scope has no such pin). Revoking
// a pin that is no longer active changes nothing.
func (d *DB) RevokeTargetPin(scope TargetScopeRef, pinID int64) (*TargetPin, error) {
if err := scope.check(); err != nil {
return nil, err
}
if _, err := d.Exec(`UPDATE target_pins SET status='revoked', revoked_at=now(), revoked_reason=$4
WHERE `+targetPinScope+` AND id=$3 AND status='active'`, scope.TaskID, scope.ConversationID, pinID, TargetPinRevokedByOperator); err != nil {
return nil, err
}
p, err := scanTargetPin(d.QueryRow(`SELECT `+targetPinCols+` FROM target_pins
WHERE `+targetPinScope+` AND id=$3`, scope.TaskID, scope.ConversationID, pinID))
if errors.Is(err, sql.ErrNoRows) {
return nil, nil
}
return p, err
}
// revokeTaskPins marks the active pins of a task that is being deleted as
// revoked. They stay as history; task ids are never reused, so a revoked pin
// can never match again.
func revokeTaskPins(tx *sql.Tx, taskID int64) error {
_, err := tx.Exec(`UPDATE target_pins SET status='revoked', revoked_at=now(), revoked_reason=$2
WHERE task_id=$1 AND status='active'`, taskID, TargetPinRevokedTaskDeleted)
return err
}
// revokeConversationPins is revokeTaskPins for deleted conversations.
func revokeConversationPins(tx *sql.Tx, convIDs []int64) error {
_, err := tx.Exec(`UPDATE target_pins SET status='revoked', revoked_at=now(), revoked_reason=$2
WHERE conversation_id=ANY($1::bigint[]) AND status='active'`, convIDs, TargetPinRevokedConversationDelete)
return err
}
// settleTaskTargetRequests closes the open target requests of a task that is
// being deleted: denied by the system, nothing pinned. Only target rows are
// touched. (A conversation needs no counterpart: its rows cascade.)
func settleTaskTargetRequests(tx *sql.Tx, taskID int64) error {
_, err := tx.Exec(`UPDATE intercept_pending SET status='denied', decided_at=NOW(), decided_by='system',
audit=CASE WHEN audit IS NULL THEN NULL ELSE audit ||
'{"effective_action":"deny","decision_reason":"task deleted","execution_status":"not_executed"}'::jsonb END
WHERE COALESCE(task_id,'')=$1 AND status='pending' AND target_key IS NOT NULL`, strconv.FormatInt(taskID, 10))
return err
}
// TargetRequest is a new target approval request: the intercept_pending row a
// gated call asks for when one of its targets has no pin that covers it.
type TargetRequest struct {
Scope TargetScopeRef
Host string
Port int
RuleID int64 // owning rule of a merged ask; 0 otherwise
AgentName string
ToolName string
ToolInput []byte
Reason string
Audit *InterceptAudit
}
// FindOrCreateTargetRequest returns the open request for (scope, host, port),
// creating it when there is none. At most one such row is open at a time
// (uq_intercept_pending_target), across goroutines and across instances that
// share the database, so concurrent calls to one host share one request.
// created reports whether this call wrote the row. id == 0 means the request
// that was open got decided in between: the caller goes back to the pin lookup.
func (d *DB) FindOrCreateTargetRequest(r TargetRequest) (id int64, created bool, err error) {
if err := r.Scope.check(); err != nil {
return 0, false, err
}
raw := json.RawMessage(r.ToolInput)
if len(raw) == 0 {
raw = json.RawMessage("{}")
}
var ruleID *int64
if r.RuleID != 0 {
ruleID = &r.RuleID
}
taskID, convID := r.Scope.pendingScope()
key := TargetKey(r.Host, r.Port)
err = d.QueryRow(`
INSERT INTO intercept_pending(rule_id, conversation_id, task_id, agent_name, tool_name, tool_input, reason, decision_source, audit, target_key)
VALUES ($1, $2, $3, $4, $5, $6, $7, 'target', $8, $9)
ON CONFLICT (COALESCE(task_id,''), COALESCE(conversation_id,0), target_key)
WHERE status = 'pending' AND target_key IS NOT NULL
DO NOTHING RETURNING id`,
ruleID, convID, taskID, r.AgentName, r.ToolName, raw, r.Reason, firstAudit([]*InterceptAudit{r.Audit}), key).Scan(&id)
if err == nil {
return id, true, nil
}
if !errors.Is(err, sql.ErrNoRows) {
return 0, false, err
}
taskKey := ""
if taskID != nil {
taskKey = *taskID
}
err = d.QueryRow(`SELECT id FROM intercept_pending
WHERE COALESCE(task_id,'')=$1 AND COALESCE(conversation_id,0)=$2 AND target_key=$3 AND status='pending'`,
taskKey, r.Scope.ConversationID, key).Scan(&id)
if errors.Is(err, sql.ErrNoRows) {
return 0, false, nil
}
return id, false, err
}
// TargetRequestKey returns the target key of an approval row: "" for a classic
// ask and for an id that does not exist.
func (d *DB) TargetRequestKey(id int64) (string, error) {
var key sql.NullString
err := d.QueryRow(`SELECT target_key FROM intercept_pending WHERE id=$1`, id).Scan(&key)
if errors.Is(err, sql.ErrNoRows) {
return "", nil
}
return key.String, err
}
// ResolveTargetRequest settles an open target request with a human decision, in
// one transaction. Allowed pins the addresses the request displayed: the active
// pin of the same key, if any, becomes superseded and a new version holds its
// addresses plus the displayed ones (addresses leave only by revocation).
// Denied pins nothing. approver and remoteAddr go to the pin and to the row's
// audit; decided_by on the row stays the actor class, human. ok is false when
// the row is not an open target request. If the pin cannot be written the row
// stays pending.
func (d *DB) ResolveTargetRequest(id int64, allowed bool, approver, remoteAddr string) (pin *TargetPin, ok bool, err error) {
tx, err := d.Begin()
if err != nil {
return nil, false, err
}
defer tx.Rollback()
var taskID sql.NullString
var convID sql.NullInt64
var raw []byte
err = tx.QueryRow(`SELECT task_id, conversation_id, audit FROM intercept_pending
WHERE id=$1 AND status='pending' AND target_key IS NOT NULL FOR UPDATE`, id).Scan(&taskID, &convID, &raw)
if errors.Is(err, sql.ErrNoRows) {
return nil, false, nil
}
if err != nil {
return nil, false, err
}
var audit InterceptAudit
if len(raw) > 0 {
if err := json.Unmarshal(raw, &audit); err != nil {
return nil, false, err
}
}
status, action, reason, execution := "denied", "deny", "user rejected the target; nothing was pinned", "not_executed"
if allowed {
if len(audit.Targets) != 1 {
return nil, false, fmt.Errorf("target request %d records %d targets, want 1", id, len(audit.Targets))
}
scope := TargetScopeRef{ConversationID: convID.Int64}
if taskID.Valid {
if scope.TaskID, err = strconv.ParseInt(taskID.String, 10, 64); err != nil {
return nil, false, fmt.Errorf("target request %d: task id %q: %w", id, taskID.String, err)
}
}
if pin, err = writeTargetPin(tx, scope, id, &audit.Targets[0], approver, remoteAddr); err != nil {
return nil, false, err
}
status, action, reason = "allowed", "allow", "user approved the target; the displayed addresses were pinned"
// A plain target request is host-level: no single call owns its result.
execution = "unknown"
if audit.MergedAsk && audit.Correlation == "exact" {
execution = "awaiting_result"
}
} else if len(audit.Targets) == 1 {
audit.Targets[0].Outcome = "denied"
}
fields := map[string]any{
"effective_action": action, "decision_reason": reason, "execution_status": execution,
"targets": audit.Targets,
}
if approver != "" {
fields["approver"] = approver
}
if remoteAddr != "" {
fields["approver_addr"] = remoteAddr
}
patch, err := json.Marshal(fields)
if err != nil {
return nil, false, err
}
if _, err := tx.Exec(`UPDATE intercept_pending SET status=$2, decided_at=NOW(), decided_by='human',
audit=COALESCE(audit,'{}'::jsonb) || $3::jsonb WHERE id=$1`, id, status, patch); err != nil {
return nil, false, err
}
if err := tx.Commit(); err != nil {
return nil, false, err
}
return pin, true, nil
}
// writeTargetPin inserts the pin an approved request creates and fills the
// request's target with what was pinned. The only writer of target_pins rows.
func writeTargetPin(tx *sql.Tx, scope TargetScopeRef, interceptID int64, t *InterceptTarget, approver, remoteAddr string) (*TargetPin, error) {
if err := scope.check(); err != nil {
return nil, err
}
// Versions count every pin the key ever had, revoked ones included, so a
// version number names one pin of the key for good.
var version int
if err := tx.QueryRow(`SELECT COALESCE(max(version),0)+1 FROM target_pins
WHERE `+targetPinScope+` AND host=$3 AND port=$4`, scope.TaskID, scope.ConversationID, t.Host, t.Port).Scan(&version); err != nil {
return nil, err
}
ips := map[string]bool{}
for _, ip := range t.ResolvedIPs {
ips[ip] = true
}
prev, err := scanTargetPin(tx.QueryRow(`SELECT `+targetPinCols+` FROM target_pins
WHERE `+targetPinScope+` AND host=$3 AND port=$4 AND status='active' FOR UPDATE`, scope.TaskID, scope.ConversationID, t.Host, t.Port))
switch {
case errors.Is(err, sql.ErrNoRows):
case err != nil:
return nil, err
default:
for _, ip := range prev.PinnedIPs {
ips[ip] = true
}
if _, err := tx.Exec(`UPDATE target_pins SET status='superseded' WHERE id=$1`, prev.ID); err != nil {
return nil, err
}
}
pinned := make([]string, 0, len(ips))
for ip := range ips {
pinned = append(pinned, ip)
}
sort.Strings(pinned)
var taskID, convID *int64
if scope.TaskID > 0 {
taskID = &scope.TaskID
} else {
convID = &scope.ConversationID
}
resolvedAt := t.ResolvedAt
if resolvedAt.IsZero() {
resolvedAt = time.Now().UTC()
}
pin, err := scanTargetPin(tx.QueryRow(`
INSERT INTO target_pins(task_id, conversation_id, host, port, version, pinned_ips, resolver, resolved_at, intercept_id, decided_by, decided_from, name_only_effective)
VALUES ($1, $2, $3, $4, $5, $6::text[], $7, $8, $9, $10, $11, $12)
RETURNING `+targetPinCols,
taskID, convID, t.Host, t.Port, version, pinned, t.Resolver, resolvedAt, interceptID, approver, remoteAddr, t.NameOnlyEffective))
if err != nil {
return nil, err
}
t.Outcome, t.PinID, t.PinVersion, t.PinnedIPs = "pinned", pin.ID, pin.Version, pin.PinnedIPs
return pin, nil
}
+924
View File
@@ -0,0 +1,924 @@
package db
import (
"bytes"
"encoding/json"
"errors"
"fmt"
"slices"
"strconv"
"strings"
"sync"
"testing"
"time"
"github.com/jackc/pgx/v5/pgconn"
)
func openTargetDB(t *testing.T) *DB {
t.Helper()
d, err := Open(testDSN(t))
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
t.Cleanup(func() { d.Close() })
return d
}
// targetTask creates a task and removes it, its pins and its approval rows at
// the end of the test (pins and approval rows outlive a task delete by design).
func targetTask(t *testing.T, d *DB, opts TaskCreateOptions) *Task {
t.Helper()
task, err := d.CreateTaskWithOptions("target gate store", "pins and requests", opts)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
_ = d.DeleteTask(task.ID)
_, _ = d.Exec(`DELETE FROM target_pins WHERE task_id=$1`, task.ID)
_, _ = d.Exec(`DELETE FROM intercept_pending WHERE task_id=$1`, strconv.FormatInt(task.ID, 10))
})
return task
}
func targetConversation(t *testing.T, d *DB) *Conversation {
t.Helper()
c, err := d.CreateConversation("assistant", fmt.Sprintf("target gate %d", time.Now().UnixNano()), nil)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
_ = d.DeleteConversation(c.ID)
_, _ = d.Exec(`DELETE FROM target_pins WHERE conversation_id=$1`, c.ID)
})
return c
}
// openTargetRequest opens the request for (scope, host, port) showing ips.
func openTargetRequest(t *testing.T, d *DB, scope TargetScopeRef, host string, port int, ips ...string) (int64, bool) {
t.Helper()
id, created, err := d.FindOrCreateTargetRequest(TargetRequest{
Scope: scope, Host: host, Port: port, AgentName: "worker", ToolName: "WebFetch",
ToolInput: []byte(`{"url":"https://` + host + `/"}`), Reason: "target " + host + " is not approved",
Audit: &InterceptAudit{
Correlation: "unavailable", InitialAction: "ask", ExecutionStatus: "not_started",
Targets: []InterceptTarget{{Host: host, Port: port, Kind: "domain", ResolvedIPs: ips, Resolver: "system",
ResolvedAt: time.Now().UTC(), Reason: "no_pin", Outcome: "requested"}},
},
})
if err != nil {
t.Fatal(err)
}
return id, created
}
// approveTarget opens a request for the key and approves it.
func approveTarget(t *testing.T, d *DB, scope TargetScopeRef, host string, port int, ips ...string) *TargetPin {
t.Helper()
id, created := openTargetRequest(t, d, scope, host, port, ips...)
if !created {
t.Fatalf("a request for %s:%d was already open", host, port)
}
pin, ok, err := d.ResolveTargetRequest(id, true, "ARTEX", "198.51.100.9")
if err != nil || !ok || pin == nil {
t.Fatalf("approve %s:%d: pin=%+v ok=%v err=%v", host, port, pin, ok, err)
}
return pin
}
func pendingTargetRows(t *testing.T, d *DB, scope TargetScopeRef) int {
t.Helper()
var n int
task := ""
if scope.TaskID > 0 {
task = strconv.FormatInt(scope.TaskID, 10)
}
if err := d.QueryRow(`SELECT count(*) FROM intercept_pending
WHERE COALESCE(task_id,'')=$1 AND COALESCE(conversation_id,0)=$2 AND status='pending' AND target_key IS NOT NULL`,
task, scope.ConversationID).Scan(&n); err != nil {
t.Fatal(err)
}
return n
}
// The whole schema is applied on every start. Applying it again must change
// nothing, add no CHECK to the gate's tables, and leave the cascade of
// intercept_pending.conversation_id as it is.
func TestTargetGateSchemaIdempotent(t *testing.T) {
snapshot := func() string {
t.Helper()
d := openTargetDB(t)
var b strings.Builder
for _, q := range []string{
`SELECT table_name||'.'||column_name||' '||data_type||' null='||is_nullable||' default='||COALESCE(column_default,'')
FROM information_schema.columns
WHERE table_schema=current_schema() AND (table_name IN ('target_pins','target_gate_exemption_log')
OR (table_name IN ('tasks','conversations') AND column_name='target_gate')
OR (table_name='intercept_pending' AND column_name='target_key'))
ORDER BY table_name, ordinal_position`,
`SELECT indexdef FROM pg_indexes WHERE schemaname=current_schema()
AND (tablename IN ('target_pins','target_gate_exemption_log') OR indexname='uq_intercept_pending_target') ORDER BY indexname`,
`SELECT conrelid::regclass::text||' '||conname||' '||contype::text||' '||confdeltype::text FROM pg_constraint
WHERE conrelid IN ('target_pins'::regclass,'target_gate_exemption_log'::regclass,'intercept_pending'::regclass)
AND contype IN ('c','f','p','u') ORDER BY 1`,
} {
rows, err := d.Query(q)
if err != nil {
t.Fatal(err)
}
for rows.Next() {
var line string
if err := rows.Scan(&line); err != nil {
t.Fatal(err)
}
b.WriteString(line + "\n")
}
if err := rows.Err(); err != nil {
t.Fatal(err)
}
rows.Close()
}
return b.String()
}
first := snapshot() // Open applied the schema
second := snapshot() // and applied it again
if first != second {
t.Fatalf("schema changed on the second application:\n--- first\n%s--- second\n%s", first, second)
}
for _, want := range []string{
"tasks.target_gate boolean null=NO default=false",
"conversations.target_gate boolean null=NO default=false",
"intercept_pending.target_key text null=YES default=",
"target_pins.name_only_effective boolean null=NO default=false",
"target_gate_exemption_log.digest text",
"uq_intercept_pending_target",
"uq_target_pins_active",
"idx_target_pins_scope",
// The cascade stays: no constraint migration.
"intercept_pending intercept_pending_conversation_id_fkey f c",
} {
if !strings.Contains(first, want) {
t.Errorf("schema snapshot lacks %q:\n%s", want, first)
}
}
// Only primary keys on the new tables: no CHECK, no foreign key.
for _, line := range strings.Split(first, "\n") {
if (strings.HasPrefix(line, "target_pins ") || strings.HasPrefix(line, "target_gate_exemption_log ")) && !strings.Contains(line, "_pkey p ") {
t.Errorf("unexpected constraint on a target-gate table: %s", line)
}
}
}
// Many callers, two instances, one key: exactly one open row, and every caller
// gets its id. This is what proves that ON CONFLICT infers the partial
// expression index.
func TestTargetRequestFindOrCreateIsIdempotent(t *testing.T) {
a, b := openTargetDB(t), openTargetDB(t)
scope := TargetScopeRef{TaskID: targetTask(t, a, TaskCreateOptions{}).ID}
const callers = 24
ids := make([]int64, callers)
made := make([]bool, callers)
var wg sync.WaitGroup
for n := range callers {
inst := a
if n%2 == 1 {
inst = b
}
wg.Go(func() {
id, created, err := inst.FindOrCreateTargetRequest(TargetRequest{
Scope: scope, Host: "app.example.com", Port: 443, ToolName: "WebFetch", Reason: "concurrent",
Audit: &InterceptAudit{Targets: []InterceptTarget{{Host: "app.example.com", Port: 443, ResolvedIPs: []string{"203.0.113.7"}}}},
})
if err != nil {
t.Error(err)
}
ids[n], made[n] = id, created
})
}
wg.Wait()
created := 0
for n := range callers {
if ids[n] == 0 || ids[n] != ids[0] {
t.Fatalf("callers got different requests: %v", ids)
}
if made[n] {
created++
}
}
if created != 1 {
t.Fatalf("%d callers created the row, want 1", created)
}
if n := pendingTargetRows(t, a, scope); n != 1 {
t.Fatalf("%d open rows, want 1", n)
}
// Another port, another scope: other keys, other rows.
if id, made := openTargetRequest(t, a, scope, "app.example.com", 80, "203.0.113.7"); !made || id == ids[0] {
t.Fatalf("port 80 shared the row of port 443 (id=%d created=%v)", id, made)
}
conv := targetConversation(t, a)
if id, made := openTargetRequest(t, a, TargetScopeRef{ConversationID: conv.ID}, "app.example.com", 443, "203.0.113.7"); !made || id == ids[0] {
t.Fatalf("a conversation shared the row of a task (id=%d created=%v)", id, made)
}
var source string
var key *string
if err := a.QueryRow(`SELECT decision_source, target_key FROM intercept_pending WHERE id=$1`, ids[0]).Scan(&source, &key); err != nil {
t.Fatal(err)
}
if source != "target" || key == nil || *key != "app.example.com:443" {
t.Fatalf("row source=%q key=%v, want target and app.example.com:443", source, key)
}
}
// A decided request does not block the key: the next call opens a new row.
// There is no sticky denial.
func TestTargetRequestNewRowAfterDecision(t *testing.T) {
d := openTargetDB(t)
scope := TargetScopeRef{TaskID: targetTask(t, d, TaskCreateOptions{}).ID}
first, _ := openTargetRequest(t, d, scope, "app.example.com", 443, "203.0.113.7")
if again, made := openTargetRequest(t, d, scope, "app.example.com", 443, "203.0.113.7"); made || again != first {
t.Fatalf("second call got id=%d created=%v, want to attach to %d", again, made, first)
}
if pin, ok, err := d.ResolveTargetRequest(first, false, "ARTEX", "198.51.100.9"); err != nil || !ok || pin != nil {
t.Fatalf("deny: pin=%+v ok=%v err=%v", pin, ok, err)
}
// A second decision on the same row is refused and changes nothing.
if _, ok, err := d.ResolveTargetRequest(first, true, "ARTEX", "198.51.100.9"); err != nil || ok {
t.Fatalf("second decision: ok=%v err=%v, want refused", ok, err)
}
if pins, err := d.ListTargetPins(scope); err != nil || len(pins) != 0 {
t.Fatalf("a denied request left pins: %+v err=%v", pins, err)
}
second, made := openTargetRequest(t, d, scope, "app.example.com", 443, "203.0.113.7")
if !made || second == first {
t.Fatalf("after a denial: id=%d created=%v, want a new row", second, made)
}
if key, err := d.TargetRequestKey(second); err != nil || key != "app.example.com:443" {
t.Fatalf("target key = %q, %v", key, err)
}
classic, err := d.CreateInterceptPending(0, 0, strconv.FormatInt(scope.TaskID, 10), "worker", "Bash", nil, "classic ask")
if err != nil {
t.Fatal(err)
}
if key, err := d.TargetRequestKey(classic); err != nil || key != "" {
t.Fatalf("classic ask has target key %q, %v", key, err)
}
// A classic ask is not a target request: it cannot be resolved as one.
if _, ok, err := d.ResolveTargetRequest(classic, true, "ARTEX", ""); err != nil || ok {
t.Fatalf("classic ask resolved as a target request: ok=%v err=%v", ok, err)
}
}
// Settling the row, superseding the old pin and writing the new one are one
// transaction: when the pin cannot be written the row stays pending and the
// previous pin stays active.
func TestResolveTargetRequestIsAtomic(t *testing.T) {
d := openTargetDB(t)
scope := TargetScopeRef{TaskID: targetTask(t, d, TaskCreateOptions{}).ID}
const host = "pin-insert-fails.example.com"
v1 := approveTarget(t, d, scope, host, 443, "203.0.113.7")
if _, err := d.Exec(`CREATE OR REPLACE FUNCTION test_refuse_target_pin() RETURNS trigger AS $$
BEGIN
IF NEW.host = '` + host + `' THEN RAISE EXCEPTION 'pin insert refused by the test'; END IF;
RETURN NEW;
END $$ LANGUAGE plpgsql`); err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
_, _ = d.Exec(`DROP TRIGGER IF EXISTS test_refuse_target_pin ON target_pins`)
_, _ = d.Exec(`DROP FUNCTION IF EXISTS test_refuse_target_pin()`)
})
if _, err := d.Exec(`CREATE TRIGGER test_refuse_target_pin BEFORE INSERT ON target_pins
FOR EACH ROW EXECUTE FUNCTION test_refuse_target_pin()`); err != nil {
t.Fatal(err)
}
id, _ := openTargetRequest(t, d, scope, host, 443, "203.0.113.99")
if pin, ok, err := d.ResolveTargetRequest(id, true, "ARTEX", "198.51.100.9"); err == nil || ok || pin != nil {
t.Fatalf("approval with a failing pin insert: pin=%+v ok=%v err=%v, want an error", pin, ok, err)
}
row, err := d.GetInterceptPending(id)
if err != nil || row == nil || row.Status != "pending" || row.DecidedBy != "" {
t.Fatalf("row after the failed approval = %+v, %v; want still pending", row, err)
}
active, err := d.ActiveTargetPins(scope, host, 443)
if err != nil || len(active) != 1 || active[0].ID != v1.ID || active[0].Version != 1 {
t.Fatalf("active pins after the failed approval = %+v, %v; want v1 untouched", active, err)
}
// Without the failure the same row is approved.
if _, err := d.Exec(`DROP TRIGGER test_refuse_target_pin ON target_pins`); err != nil {
t.Fatal(err)
}
if pin, ok, err := d.ResolveTargetRequest(id, true, "ARTEX", "198.51.100.9"); err != nil || !ok || pin == nil || pin.Version != 2 {
t.Fatalf("approval after the failure was removed: pin=%+v ok=%v err=%v", pin, ok, err)
}
}
// Approving a request for a pinned key extends the pin: the new version holds
// the old addresses plus the displayed ones. Addresses leave only by revocation.
func TestTargetPinVersioningExtends(t *testing.T) {
d := openTargetDB(t)
scope := TargetScopeRef{TaskID: targetTask(t, d, TaskCreateOptions{}).ID}
v1 := approveTarget(t, d, scope, "app.example.com", 443, "203.0.113.7")
if v1.Version != 1 || !slices.Equal(v1.PinnedIPs, []string{"203.0.113.7"}) || v1.Status != TargetPinActive ||
v1.DecidedBy != "ARTEX" || v1.DecidedFrom != "198.51.100.9" || v1.Resolver != "system" || v1.InterceptID == nil ||
v1.TaskID == nil || *v1.TaskID != scope.TaskID || v1.ConversationID != nil || v1.NameOnlyEffective {
t.Fatalf("v1 = %+v", v1)
}
v2 := approveTarget(t, d, scope, "app.example.com", 443, "203.0.113.9", "203.0.113.7", "203.0.113.8")
if v2.Version != 2 || !slices.Equal(v2.PinnedIPs, []string{"203.0.113.7", "203.0.113.8", "203.0.113.9"}) {
t.Fatalf("v2 = %+v, want version 2 with the union, sorted", v2)
}
pins, err := d.ListTargetPins(scope)
if err != nil || len(pins) != 2 || pins[0].ID != v2.ID || pins[0].Status != TargetPinActive ||
pins[1].ID != v1.ID || pins[1].Status != TargetPinSuperseded {
t.Fatalf("pins = %+v, %v; want v2 active then v1 superseded", pins, err)
}
// The request row records what was pinned and by whom; decided_by on the row
// stays the actor class.
detail, err := d.GetInterceptDetail(*v2.InterceptID)
if err != nil || detail == nil || detail.Status != "allowed" || detail.DecidedBy != "human" || detail.DecisionSource != "target" ||
detail.Audit == nil || detail.Audit.Approver != "ARTEX" || detail.Audit.ApproverAddr != "198.51.100.9" ||
detail.Audit.EffectiveAction != "allow" || len(detail.Audit.Targets) != 1 {
t.Fatalf("request row = %+v, %v", detail, err)
}
if got := detail.Audit.Targets[0]; got.Outcome != "pinned" || got.PinID != v2.ID || got.PinVersion != 2 || len(got.PinnedIPs) != 3 {
t.Fatalf("recorded target = %+v", got)
}
}
// One active pin per key is a database fact, across instances.
func TestTargetPinOneActivePerKey(t *testing.T) {
a, b := openTargetDB(t), openTargetDB(t)
scope := TargetScopeRef{TaskID: targetTask(t, a, TaskCreateOptions{}).ID}
approveTarget(t, a, scope, "app.example.com", 443, "203.0.113.7")
insert := func(port int, status string) error {
_, err := b.Exec(`INSERT INTO target_pins(task_id, host, port, version, pinned_ips, status)
VALUES ($1, 'app.example.com', $2, 9, '{203.0.113.99}', $3)`, scope.TaskID, port, status)
return err
}
var pgErr *pgconn.PgError
if err := insert(443, "active"); !errors.As(err, &pgErr) || pgErr.Code != "23505" {
t.Fatalf("second active pin on one key: %v, want a unique violation", err)
}
if err := insert(443, "superseded"); err != nil {
t.Fatalf("a non-active pin on the same key must be allowed: %v", err)
}
if err := insert(8443, "active"); err != nil {
t.Fatalf("an active pin on another port must be allowed: %v", err)
}
approveTarget(t, b, scope, "app.example.com", 443, "203.0.113.8")
var active int
if err := a.QueryRow(`SELECT count(*) FROM target_pins WHERE task_id=$1 AND host='app.example.com' AND port=443 AND status='active'`,
scope.TaskID).Scan(&active); err != nil || active != 1 {
t.Fatalf("%d active pins on the key after a second approval (%v), want 1", active, err)
}
}
func TestTargetPinRevokeIdempotent(t *testing.T) {
d := openTargetDB(t)
scope := TargetScopeRef{TaskID: targetTask(t, d, TaskCreateOptions{}).ID}
other := TargetScopeRef{TaskID: targetTask(t, d, TaskCreateOptions{}).ID}
pin := approveTarget(t, d, scope, "app.example.com", 443, "203.0.113.7")
// Another scope cannot see or revoke the pin.
if got, err := d.RevokeTargetPin(other, pin.ID); err != nil || got != nil {
t.Fatalf("revoke through another scope = %+v, %v; want nil", got, err)
}
if got, err := d.RevokeTargetPin(scope, pin.ID+1000000); err != nil || got != nil {
t.Fatalf("revoke of an unknown pin = %+v, %v; want nil", got, err)
}
if active, err := d.ActiveTargetPins(scope, "app.example.com", 443); err != nil || len(active) != 1 {
t.Fatalf("pin no longer active after foreign revokes: %+v, %v", active, err)
}
first, err := d.RevokeTargetPin(scope, pin.ID)
if err != nil || first == nil || first.Status != TargetPinRevoked || first.RevokedReason != TargetPinRevokedByOperator || first.RevokedAt == nil {
t.Fatalf("revoke = %+v, %v", first, err)
}
second, err := d.RevokeTargetPin(scope, pin.ID)
if err != nil || second == nil || second.Status != TargetPinRevoked || !second.RevokedAt.Equal(*first.RevokedAt) {
t.Fatalf("second revoke = %+v, %v; want the first revocation unchanged", second, err)
}
if active, err := d.ActiveTargetPins(scope, "app.example.com", 443); err != nil || len(active) != 0 {
t.Fatalf("a revoked pin still matches: %+v, %v", active, err)
}
// A new approval starts from the displayed addresses only, under a new version.
again := approveTarget(t, d, scope, "app.example.com", 443, "203.0.113.50")
if again.Version != 2 || !slices.Equal(again.PinnedIPs, []string{"203.0.113.50"}) {
t.Fatalf("pin after a revocation = %+v, want version 2 holding only the new address", again)
}
}
// A pin on port 0 (port not declared) covers every port of the host. A pin on
// a specific port covers that port only, and never a target whose port is not
// declared.
func TestActiveTargetPinPortPrecedence(t *testing.T) {
d := openTargetDB(t)
scope := TargetScopeRef{TaskID: targetTask(t, d, TaskCreateOptions{}).ID}
anyPort := approveTarget(t, d, scope, "both.example.com", 0, "203.0.113.1")
exact := approveTarget(t, d, scope, "both.example.com", 443, "203.0.113.2")
only443 := approveTarget(t, d, scope, "exact.example.com", 443, "203.0.113.3")
for _, tc := range []struct {
host string
port int
want []int64
}{
{"both.example.com", 443, []int64{exact.ID, anyPort.ID}}, // the exact port first
{"both.example.com", 80, []int64{anyPort.ID}},
{"both.example.com", 8443, []int64{anyPort.ID}},
{"both.example.com", 0, []int64{anyPort.ID}}, // never the pin on 443
{"exact.example.com", 443, []int64{only443.ID}},
{"exact.example.com", 80, nil},
{"exact.example.com", 0, nil}, // a specific port never covers "not declared"
{"unknown.example.com", 443, nil},
} {
pins, err := d.ActiveTargetPins(scope, tc.host, tc.port)
if err != nil {
t.Fatal(err)
}
var got []int64
for _, p := range pins {
got = append(got, p.ID)
}
if !slices.Equal(got, tc.want) {
t.Errorf("candidates for %s:%d = %v, want %v", tc.host, tc.port, got, tc.want)
}
}
// The address sets stay apart: each candidate holds its own.
pins, _ := d.ActiveTargetPins(scope, "both.example.com", 443)
if len(pins) != 2 || !slices.Equal(pins[0].PinnedIPs, []string{"203.0.113.2"}) || !slices.Equal(pins[1].PinnedIPs, []string{"203.0.113.1"}) {
t.Fatalf("candidates = %+v, want two pins with separate addresses", pins)
}
}
// A hard delete keeps the gate's history: pins revoked, approval rows readable,
// no target request left open.
func TestTaskDeleteKeepsPinsAndApprovalRows(t *testing.T) {
d := openTargetDB(t)
task := targetTask(t, d, TaskCreateOptions{})
other := targetTask(t, d, TaskCreateOptions{})
scope := TargetScopeRef{TaskID: task.ID}
pin := approveTarget(t, d, scope, "app.example.com", 443, "203.0.113.7")
kept := approveTarget(t, d, TargetScopeRef{TaskID: other.ID}, "app.example.com", 443, "203.0.113.7")
open, _ := openTargetRequest(t, d, scope, "new.example.com", 443, "203.0.113.20")
classic, err := d.CreateInterceptPending(0, 0, strconv.FormatInt(task.ID, 10), "worker", "Bash", nil, "classic ask")
if err != nil {
t.Fatal(err)
}
if err := d.DeleteTask(task.ID); err != nil {
t.Fatal(err)
}
pins, err := d.ListTargetPins(scope)
if err != nil || len(pins) != 1 || pins[0].ID != pin.ID || pins[0].Status != TargetPinRevoked ||
pins[0].RevokedReason != TargetPinRevokedTaskDeleted || pins[0].RevokedAt == nil ||
!slices.Equal(pins[0].PinnedIPs, []string{"203.0.113.7"}) || pins[0].DecidedBy != "ARTEX" {
t.Fatalf("pins of the deleted task = %+v, %v; want one revoked pin, still readable", pins, err)
}
if active, err := d.ActiveTargetPins(scope, "app.example.com", 443); err != nil || len(active) != 0 {
t.Fatalf("a pin of the deleted task still matches: %+v, %v", active, err)
}
if active, err := d.ActiveTargetPins(TargetScopeRef{TaskID: other.ID}, "app.example.com", 443); err != nil || len(active) != 1 || active[0].ID != kept.ID {
t.Fatalf("the pin of another task was touched: %+v, %v", active, err)
}
approved, err := d.GetInterceptDetail(*pin.InterceptID)
if err != nil || approved == nil || approved.Status != "allowed" || approved.Audit == nil || approved.Audit.Approver != "ARTEX" {
t.Fatalf("approval row of the deleted task = %+v, %v; want still readable", approved, err)
}
settled, err := d.GetInterceptDetail(open)
if err != nil || settled == nil || settled.Status != "denied" || settled.DecidedBy != "system" ||
settled.Audit == nil || settled.Audit.EffectiveAction != "deny" || settled.Audit.DecisionReason != "task deleted" {
t.Fatalf("open target request after the delete = %+v, %v; want denied by the system", settled, err)
}
if n := pendingTargetRows(t, d, scope); n != 0 {
t.Fatalf("%d target requests still open after the delete", n)
}
// Classic asks are not the gate's: untouched, as before.
if row, err := d.GetInterceptPending(classic); err != nil || row == nil || row.Status != "pending" {
t.Fatalf("classic ask after the delete = %+v, %v; want untouched", row, err)
}
}
// archiveTargetTask moves a task into the cold archive and returns the archive
// id and its snapshot.
func archiveTargetTask(t *testing.T, d *DB, taskID int64) (int64, *TaskArchiveSnapshot, *bytes.Buffer) {
t.Helper()
if err := d.EnsureLLMRecordsTable(); err != nil {
t.Fatal(err)
}
if err := d.EnsureLLMUsageTable(); err != nil {
t.Fatal(err)
}
if err := d.SetPaused(taskID, true); err != nil {
t.Fatal(err)
}
archive, err := d.QueueTaskArchive(taskID)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _, _ = d.Exec(`DELETE FROM task_archives WHERE id=$1`, archive.ID) })
if claimed, err := d.ClaimTaskArchiveJob(t.Context()); err != nil || claimed == nil || claimed.ID != archive.ID {
t.Fatalf("archive claim = %+v, %v", claimed, err)
}
var llmRecords bytes.Buffer
snapshot, err := d.SnapshotTaskArchiveWithLLMRecords(taskID, &llmRecords)
if err != nil {
t.Fatal(err)
}
if err := d.CompleteTaskArchive(archive.ID, snapshot, "/tmp/target-gate-task.tar.zst", "abc", 100, 50); err != nil {
t.Fatal(err)
}
return archive.ID, snapshot, &llmRecords
}
func TestArchiveStubDeleteRevokesPins(t *testing.T) {
d := openTargetDB(t)
task := targetTask(t, d, TaskCreateOptions{})
scope := TargetScopeRef{TaskID: task.ID}
pin := approveTarget(t, d, scope, "app.example.com", 443, "203.0.113.7")
archiveID, _, _ := archiveTargetTask(t, d, task.ID)
if active, err := d.ActiveTargetPins(scope, "app.example.com", 443); err != nil || len(active) != 1 {
t.Fatalf("pin while the task is an archived stub = %+v, %v; want still active", active, err)
}
if _, err := d.QueueTaskArchiveDelete(archiveID); err != nil {
t.Fatal(err)
}
if claimed, err := d.ClaimTaskArchiveJob(t.Context()); err != nil || claimed == nil || claimed.State != Deleting {
t.Fatalf("delete claim = %+v, %v", claimed, err)
}
if err := d.DeleteTaskArchiveStub(archiveID); err != nil {
t.Fatal(err)
}
pins, err := d.ListTargetPins(scope)
if err != nil || len(pins) != 1 || pins[0].ID != pin.ID || pins[0].Status != TargetPinRevoked || pins[0].RevokedReason != TargetPinRevokedTaskDeleted {
t.Fatalf("pins after the stub delete = %+v, %v; want one pin revoked as task_deleted", pins, err)
}
}
// The archive carries neither pins nor the flag: both stay in the live tables
// while the task is a stub and are there after the restore.
func TestArchiveRestoreKeepsPinsAndFlag(t *testing.T) {
d := openTargetDB(t)
on := true
task := targetTask(t, d, TaskCreateOptions{TargetGate: &on})
scope := TargetScopeRef{TaskID: task.ID}
pin := approveTarget(t, d, scope, "app.example.com", 443, "203.0.113.7")
open, _ := openTargetRequest(t, d, scope, "new.example.com", 443, "203.0.113.20")
archiveID, snapshot, llmRecords := archiveTargetTask(t, d, task.ID)
var flag bool
if err := d.QueryRow(`SELECT target_gate FROM tasks WHERE id=$1`, task.ID).Scan(&flag); err != nil || !flag {
t.Fatalf("flag of the archived stub = %v, %v; want kept", flag, err)
}
if _, err := d.QueueTaskArchiveRestore(archiveID); err != nil {
t.Fatal(err)
}
if claimed, err := d.ClaimTaskArchiveJob(t.Context()); err != nil || claimed == nil || claimed.State != Restoring {
t.Fatalf("restore claim = %+v, %v", claimed, err)
}
if _, err := d.RestoreTaskArchiveWithLLMRecords(archiveID, snapshot, 0, bytes.NewReader(llmRecords.Bytes())); err != nil {
t.Fatal(err)
}
if err := d.CompleteTaskArchiveRestore(archiveID); err != nil {
t.Fatal(err)
}
live, err := d.GetTask(task.ID)
if err != nil || live == nil || !live.TargetGate {
t.Fatalf("restored task = %+v, %v; want the gate still on", live, err)
}
active, err := d.ActiveTargetPins(scope, "app.example.com", 443)
if err != nil || len(active) != 1 || active[0].ID != pin.ID {
t.Fatalf("pins after the restore = %+v, %v; want the pin still active", active, err)
}
// The request that was open at archive time comes back closed (a timeout,
// never an approval) with its key, and the key is free again.
row, err := d.GetInterceptPending(open)
if err != nil || row == nil || row.Status != "timeout" || row.DecidedBy != "system" {
t.Fatalf("restored open request = %+v, %v; want a system timeout", row, err)
}
if key, err := d.TargetRequestKey(open); err != nil || key != "new.example.com:443" {
t.Fatalf("restored target key = %q, %v", key, err)
}
if id, made := openTargetRequest(t, d, scope, "new.example.com", 443, "203.0.113.20"); !made || id == open {
t.Fatalf("request after the restore: id=%d created=%v, want a new row", id, made)
}
}
// Archives written before the gate have no target_key: the row restores with
// NULL, a classic ask.
func TestRestoreInterceptRowsWithoutTargetKey(t *testing.T) {
d := openTargetDB(t)
task := targetTask(t, d, TaskCreateOptions{})
id, err := d.CreateDecidedIntercept(0, 0, strconv.FormatInt(task.ID, 10), "worker", "Bash", nil, "allowed", "[model] allow",
&InterceptAudit{InitialAction: "allow", ExecutionStatus: "succeeded"})
if err != nil {
t.Fatal(err)
}
var raw []byte
if err := d.QueryRow(`SELECT row_to_json(ip) FROM intercept_pending ip WHERE id=$1`, id).Scan(&raw); err != nil {
t.Fatal(err)
}
var row map[string]any
if err := json.Unmarshal(raw, &row); err != nil {
t.Fatal(err)
}
if _, ok := row["target_key"]; !ok {
t.Fatal("control: the live row was expected to carry the target_key column")
}
delete(row, "target_key")
archived, _ := json.Marshal([]map[string]any{row})
tx, err := d.Begin()
if err != nil {
t.Fatal(err)
}
defer tx.Rollback()
if _, err := tx.Exec(`DELETE FROM intercept_pending WHERE id=$1`, id); err != nil {
t.Fatal(err)
}
if err := restoreInterceptRows(tx, archived); err != nil {
t.Fatal(err)
}
var key *string
var status string
if err := tx.QueryRow(`SELECT target_key, status FROM intercept_pending WHERE id=$1`, id).Scan(&key, &status); err != nil {
t.Fatal(err)
}
if key != nil || status != "allowed" {
t.Fatalf("restored row target_key=%v status=%s, want NULL and allowed", key, status)
}
}
// Deleting a conversation revokes its pins in the same transaction. The pins
// stay readable; the request rows go with the conversation (cascade). A delete
// that fails leaves the pins active.
func TestConversationDeleteRevokesPinsInSameTx(t *testing.T) {
d := openTargetDB(t)
conv, other := targetConversation(t, d), targetConversation(t, d)
scope := TargetScopeRef{ConversationID: conv.ID}
pin := approveTarget(t, d, scope, "app.example.com", 443, "203.0.113.7")
kept := approveTarget(t, d, TargetScopeRef{ConversationID: other.ID}, "app.example.com", 443, "203.0.113.7")
openTargetRequest(t, d, scope, "new.example.com", 443, "203.0.113.20")
if pin.ConversationID == nil || *pin.ConversationID != conv.ID || pin.TaskID != nil {
t.Fatalf("conversation pin = %+v", pin)
}
// A reference without a delete rule makes the delete fail after the revoke
// statement has run inside the transaction.
if _, err := d.Exec(`CREATE TABLE IF NOT EXISTS test_conversation_delete_block (conversation_id BIGINT REFERENCES conversations(id))`); err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _, _ = d.Exec(`DROP TABLE IF EXISTS test_conversation_delete_block`) })
if _, err := d.Exec(`INSERT INTO test_conversation_delete_block VALUES ($1)`, conv.ID); err != nil {
t.Fatal(err)
}
if err := d.DeleteConversation(conv.ID); err == nil {
t.Fatal("control: the delete was expected to fail on the blocking reference")
}
if active, err := d.ActiveTargetPins(scope, "app.example.com", 443); err != nil || len(active) != 1 || active[0].ID != pin.ID {
t.Fatalf("pins after a failed delete = %+v, %v; want the pin still active", active, err)
}
if got, err := d.GetConversation(conv.ID); err != nil || got == nil {
t.Fatalf("conversation after a failed delete = %+v, %v", got, err)
}
if _, err := d.Exec(`DELETE FROM test_conversation_delete_block`); err != nil {
t.Fatal(err)
}
if err := d.DeleteConversation(conv.ID); err != nil {
t.Fatal(err)
}
pins, err := d.ListTargetPins(scope)
if err != nil || len(pins) != 1 || pins[0].ID != pin.ID || pins[0].Status != TargetPinRevoked ||
pins[0].RevokedReason != TargetPinRevokedConversationDelete || pins[0].RevokedAt == nil ||
!slices.Equal(pins[0].PinnedIPs, []string{"203.0.113.7"}) || pins[0].DecidedBy != "ARTEX" || pins[0].InterceptID == nil {
t.Fatalf("pins of the deleted conversation = %+v, %v; want one revoked pin, still readable", pins, err)
}
var rows int
if err := d.QueryRow(`SELECT count(*) FROM intercept_pending WHERE conversation_id=$1`, conv.ID).Scan(&rows); err != nil || rows != 0 {
t.Fatalf("%d request rows left after the delete (%v), want 0: the cascade stays", rows, err)
}
// The pin names a request row that is gone; it is kept as a plain number.
if gone, err := d.GetInterceptPending(*pins[0].InterceptID); err != nil || gone != nil {
t.Fatalf("request row of the pin = %+v, %v; want gone", gone, err)
}
if active, err := d.ActiveTargetPins(TargetScopeRef{ConversationID: other.ID}, "app.example.com", 443); err != nil || len(active) != 1 || active[0].ID != kept.ID {
t.Fatalf("the pin of another conversation was touched: %+v, %v", active, err)
}
}
// The batch delete revokes the pins of the ids it actually removed, and only
// those.
func TestConversationBatchDeleteRevokesOnlyDeletedIds(t *testing.T) {
d := openTargetDB(t)
gone, stays := targetConversation(t, d), targetConversation(t, d)
approveTarget(t, d, TargetScopeRef{ConversationID: gone.ID}, "app.example.com", 443, "203.0.113.7")
approveTarget(t, d, TargetScopeRef{ConversationID: stays.ID}, "app.example.com", 443, "203.0.113.7")
// A pin whose conversation id is in the request but does not exist: the
// delete removes nothing for it, so the pin must not be revoked.
const unknown = int64(1) << 60
if _, err := d.Exec(`INSERT INTO target_pins(conversation_id, host, port, pinned_ips) VALUES ($1, 'app.example.com', 443, '{203.0.113.7}')`, unknown); err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _, _ = d.Exec(`DELETE FROM target_pins WHERE conversation_id=$1`, unknown) })
deleted, err := d.DeleteConversations([]int64{gone.ID, unknown})
if err != nil || !slices.Equal(deleted, []int64{gone.ID}) {
t.Fatalf("deleted = %v, %v; want only %d", deleted, err, gone.ID)
}
status := func(convID int64) (string, string) {
t.Helper()
pins, err := d.ListTargetPins(TargetScopeRef{ConversationID: convID})
if err != nil || len(pins) != 1 {
t.Fatalf("pins of conversation %d = %+v, %v", convID, pins, err)
}
return pins[0].Status, pins[0].RevokedReason
}
if s, reason := status(gone.ID); s != TargetPinRevoked || reason != TargetPinRevokedConversationDelete {
t.Fatalf("pin of the deleted conversation: %s/%s, want revoked/conversation_deleted", s, reason)
}
if s, _ := status(stays.ID); s != TargetPinActive {
t.Fatalf("pin of a conversation that was not in the request: %s, want active", s)
}
if s, _ := status(unknown); s != TargetPinActive {
t.Fatalf("pin of an id the delete did not remove: %s, want active", s)
}
}
func TestTaskTargetGateDefaultsOff(t *testing.T) {
d := openTargetDB(t)
task := targetTask(t, d, TaskCreateOptions{})
if task.TargetGate {
t.Fatal("a task created without the option has the gate on")
}
got, err := d.GetTask(task.ID)
if err != nil || got == nil || got.TargetGate {
t.Fatalf("stored task = %+v, %v; want the gate off", got, err)
}
off := false
if explicit := targetTask(t, d, TaskCreateOptions{TargetGate: &off}); explicit.TargetGate {
t.Fatal("a task created with the option false has the gate on")
}
c := targetConversation(t, d)
if c.TargetGate {
t.Fatal("a conversation created without the flag has the gate on")
}
}
func TestTaskTargetGatePatchRoundTrip(t *testing.T) {
d := openTargetDB(t)
on, off := true, false
task := targetTask(t, d, TaskCreateOptions{TargetGate: &on})
if !task.TargetGate {
t.Fatal("create with the option true returned the gate off")
}
if got, err := d.GetTask(task.ID); err != nil || got == nil || !got.TargetGate {
t.Fatalf("stored task = %+v, %v; want the gate on", got, err)
}
// A patch of another field leaves the flag alone.
name := "renamed"
if got, err := d.UpdateTask(task.ID, TaskPatch{Name: &name}); err != nil || got == nil || !got.TargetGate || got.Name != name {
t.Fatalf("after a name patch = %+v, %v; want the gate still on", got, err)
}
if got, err := d.UpdateTask(task.ID, TaskPatch{TargetGate: &off}); err != nil || got == nil || got.TargetGate || got.Name != name {
t.Fatalf("after switching off = %+v, %v", got, err)
}
if got, err := d.UpdateTask(task.ID, TaskPatch{TargetGate: &on}); err != nil || got == nil || !got.TargetGate {
t.Fatalf("after switching on = %+v, %v", got, err)
}
listed, err := d.ListTasks()
if err != nil {
t.Fatal(err)
}
for _, l := range listed {
if l.ID == task.ID && !l.TargetGate {
t.Fatal("the task list reports the gate off")
}
}
}
func TestConversationTargetGateRoundTrip(t *testing.T) {
d := openTargetDB(t)
c, err := d.CreateConversationGated("assistant", fmt.Sprintf("gated %d", time.Now().UnixNano()), nil, true)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = d.DeleteConversation(c.ID) })
if !c.TargetGate {
t.Fatal("create with the flag returned the gate off")
}
if got, err := d.GetConversation(c.ID); err != nil || got == nil || !got.TargetGate {
t.Fatalf("stored conversation = %+v, %v; want the gate on", got, err)
}
title, on, off := "renamed", true, false
if got, err := d.UpdateConversation(c.ID, ConversationPatch{Title: &title}); err != nil || got == nil || !got.TargetGate || got.Title != title {
t.Fatalf("after a title patch = %+v, %v; want the gate still on", got, err)
}
if got, err := d.UpdateConversation(c.ID, ConversationPatch{TargetGate: &off}); err != nil || got == nil || got.TargetGate || got.Title != title {
t.Fatalf("after switching off = %+v, %v", got, err)
}
if got, err := d.UpdateConversation(c.ID, ConversationPatch{TargetGate: &on}); err != nil || got == nil || !got.TargetGate {
t.Fatalf("after switching on = %+v, %v", got, err)
}
listed, err := d.ListConversations()
if err != nil {
t.Fatal(err)
}
for _, l := range listed {
if l.ID == c.ID && !l.TargetGate {
t.Fatal("the conversation list reports the gate off")
}
}
}
// A retest conversation takes the flag of the finding's task and none of its
// pins.
func TestRetestConversationCopiesFlagNotPins(t *testing.T) {
d := openTargetDB(t)
on := true
for _, tc := range []struct {
name string
opts TaskCreateOptions
want bool
}{
{"gate on", TaskCreateOptions{TargetGate: &on}, true},
{"gate off", TaskCreateOptions{}, false},
} {
t.Run(tc.name, func(t *testing.T) {
task := targetTask(t, d, tc.opts)
approveTarget(t, d, TargetScopeRef{TaskID: task.ID}, "app.example.com", 443, "203.0.113.7")
fid, err := d.AddFinding(task.ID, 0, "retest-gate", "gated finding", "high", "summary", "evidence", "worker", nil)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
_, _ = d.Exec(`DELETE FROM conversations WHERE id IN (SELECT conversation_id FROM finding_retests WHERE finding_id=$1)`, fid)
_, _ = d.DeleteFinding(fid)
})
_, conv, created, err := d.CreateFindingRetest(t.Context(), fid, "")
if err != nil || !created || conv == nil {
t.Fatalf("retest: conv=%+v created=%v err=%v", conv, created, err)
}
if conv.TargetGate != tc.want {
t.Fatalf("retest conversation gate = %v, want %v", conv.TargetGate, tc.want)
}
if stored, err := d.GetConversation(conv.ID); err != nil || stored == nil || stored.TargetGate != tc.want {
t.Fatalf("stored retest conversation = %+v, %v", stored, err)
}
if pins, err := d.ListTargetPins(TargetScopeRef{ConversationID: conv.ID}); err != nil || len(pins) != 0 {
t.Fatalf("retest conversation inherited pins: %+v, %v", pins, err)
}
})
}
// A finding without a task: the gate is off.
fid, err := d.AddFinding(0, 0, "retest-gate", "orphan finding", "high", "summary", "evidence", "worker", nil)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
_, _ = d.Exec(`DELETE FROM conversations WHERE id IN (SELECT conversation_id FROM finding_retests WHERE finding_id=$1)`, fid)
_, _ = d.DeleteFinding(fid)
})
if _, conv, _, err := d.CreateFindingRetest(t.Context(), fid, ""); err != nil || conv == nil || conv.TargetGate {
t.Fatalf("retest of a finding without a task: conv=%+v err=%v; want the gate off", conv, err)
}
}
// Every human decision records who took it and from where, classic asks
// included. decided_by keeps the actor class.
func TestResolveInterceptByRecordsApprover(t *testing.T) {
d := openTargetDB(t)
task := targetTask(t, d, TaskCreateOptions{})
create := func() int64 {
t.Helper()
id, err := d.CreateInterceptPending(0, 0, strconv.FormatInt(task.ID, 10), "worker", "Bash", nil, "classic ask",
&InterceptAudit{InitialAction: "ask", ExecutionStatus: "not_started", Correlation: "exact"})
if err != nil {
t.Fatal(err)
}
return id
}
id := create()
if ok, err := d.ResolveInterceptBy(id, "allowed", "allow", "user approved execution", "human", "ARTEX", "198.51.100.9"); err != nil || !ok {
t.Fatalf("resolve: ok=%v err=%v", ok, err)
}
got, err := d.GetInterceptDetail(id)
if err != nil || got == nil || got.DecidedBy != "human" || got.Audit == nil || got.Audit.Approver != "ARTEX" ||
got.Audit.ApproverAddr != "198.51.100.9" || got.Audit.EffectiveAction != "allow" || got.Audit.ExecutionStatus != "awaiting_result" {
t.Fatalf("decided classic ask = %+v, %v", got, err)
}
// A system decision has no author.
id = create()
if ok, err := d.ResolveIntercept(id, "timeout", "deny", "approval timed out", "system"); err != nil || !ok {
t.Fatalf("resolve: ok=%v err=%v", ok, err)
}
var raw []byte
if err := d.QueryRow(`SELECT audit FROM intercept_pending WHERE id=$1`, id).Scan(&raw); err != nil {
t.Fatal(err)
}
if strings.Contains(string(raw), "approver") {
t.Fatalf("a system decision recorded an approver: %s", raw)
}
}
+7
View File
@@ -766,6 +766,13 @@ WHERE archive.id=$1 FOR UPDATE OF archive,task`, archiveID).Scan(&taskID, &expID
if !errors.Is(err, sql.ErrNoRows) {
return err
}
// As in DeleteTaskCascadePrepared: pins stay as revoked history.
if err := revokeTaskPins(tx, taskID); err != nil {
return err
}
if err := settleTaskTargetRequests(tx, taskID); err != nil {
return err
}
if _, err := tx.Exec(`DELETE FROM tasks WHERE id=$1 AND archived_at IS NOT NULL`, taskID); err != nil {
return err
}
+24 -7
View File
@@ -52,6 +52,9 @@ type Task struct {
// PlanQuietMode gates the planner quiet-mode (#86): edge-triggered rounds
// are skipped while work is in flight and the open-goal set is unchanged.
PlanQuietMode bool `json:"plan_quiet_mode"`
// TargetGate is the per-task opt-in of the approval gate for asset targets.
// false (default) = no gate, behaviour unchanged.
TargetGate bool `json:"target_gate"`
// CoverageEnabled is the master switch of the "asset coverage" feature (default true). When false:
// do not compute/show test coverage, do not auto-accumulate task_scope (source=auto), do not open
// add_task_scope/list_untested_assets to the agent, and do not inject the coverage block into the
@@ -173,6 +176,8 @@ type TaskCreateOptions struct {
// chain instead of the single active head. nil = off (default), preserving
// the ordered-failover semantics.
LLMRotate *bool
// TargetGate opts the task into the approval gate for asset targets; nil = off.
TargetGate *bool
}
// CreateTaskWithOptions creates an exploration, task, direct source relations,
@@ -251,10 +256,11 @@ VALUES ($1, 'fact', $2, 0, 'origin', 'system')`, expID, string(originPayload));
if opts.LLMRotate != nil {
rotate = *opts.LLMRotate
}
t.TargetGate = opts.TargetGate != nil && *opts.TargetGate
if err := tx.QueryRow(`
INSERT INTO tasks(name, category_id, description, goal, exploration_id, llm_profile_id, active_llm_profile_id, timeout_seconds, plan_heartbeat_seconds, plan_quiet_mode, coverage_enabled, llm_rotate)
VALUES ($1,$2,$3,$4,$5,$6,$6,$7,$8,$9,$10,$11)
RETURNING id, status, paused, created_at`, opts.Name, opts.CategoryID, description, goal, expID, active, opts.TimeoutSeconds, opts.PlanHeartbeatSeconds, quietMode, coverageEnabled, rotate).Scan(&t.ID, &t.Status, &t.Paused, &t.CreatedAt); err != nil {
INSERT INTO tasks(name, category_id, description, goal, exploration_id, llm_profile_id, active_llm_profile_id, timeout_seconds, plan_heartbeat_seconds, plan_quiet_mode, coverage_enabled, llm_rotate, target_gate)
VALUES ($1,$2,$3,$4,$5,$6,$6,$7,$8,$9,$10,$11,$12)
RETURNING id, status, paused, created_at`, opts.Name, opts.CategoryID, description, goal, expID, active, opts.TimeoutSeconds, opts.PlanHeartbeatSeconds, quietMode, coverageEnabled, rotate, t.TargetGate).Scan(&t.ID, &t.Status, &t.Paused, &t.CreatedAt); err != nil {
return nil, err
}
if err := insertTaskRelations(tx, t.ID, opts.SourceTaskIDs); err != nil {
@@ -370,11 +376,11 @@ func insertTaskLLMProfiles(tx *sql.Tx, taskID int64, profileIDs []int64) error {
const taskCols = `id, COALESCE(name,''), category_id,
COALESCE((SELECT category.name FROM task_categories category WHERE category.id=tasks.category_id),''),
description, goal, exploration_id, status, paused, queued, queued_at, COALESCE(queue_mode,''), llm_profile_id, active_llm_profile_id, COALESCE(parent_ref,''), pinned_at, created_at, completed_at, COALESCE(timeout_seconds,0), COALESCE(plan_heartbeat_seconds,300), COALESCE(plan_quiet_mode,true), COALESCE(coverage_enabled,true), COALESCE(llm_rotate,false), first_run_at, deadline_at`
description, goal, exploration_id, status, paused, queued, queued_at, COALESCE(queue_mode,''), llm_profile_id, active_llm_profile_id, COALESCE(parent_ref,''), pinned_at, created_at, completed_at, COALESCE(timeout_seconds,0), COALESCE(plan_heartbeat_seconds,300), COALESCE(plan_quiet_mode,true), COALESCE(coverage_enabled,true), COALESCE(llm_rotate,false), first_run_at, deadline_at, target_gate`
func scanTask(sc interface{ Scan(...any) error }) (*Task, error) {
var t Task
if err := sc.Scan(&t.ID, &t.Name, &t.CategoryID, &t.CategoryName, &t.Description, &t.Goal, &t.ExplorationID, &t.Status, &t.Paused, &t.Queued, &t.QueuedAt, &t.QueueMode, &t.LLMProfileID, &t.ActiveLLMProfileID, &t.ParentRef, &t.PinnedAt, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.PlanHeartbeatSeconds, &t.PlanQuietMode, &t.CoverageEnabled, &t.LLMRotate, &t.FirstRunAt, &t.DeadlineAt); err != nil {
if err := sc.Scan(&t.ID, &t.Name, &t.CategoryID, &t.CategoryName, &t.Description, &t.Goal, &t.ExplorationID, &t.Status, &t.Paused, &t.Queued, &t.QueuedAt, &t.QueueMode, &t.LLMProfileID, &t.ActiveLLMProfileID, &t.ParentRef, &t.PinnedAt, &t.CreatedAt, &t.CompletedAt, &t.TimeoutSeconds, &t.PlanHeartbeatSeconds, &t.PlanQuietMode, &t.CoverageEnabled, &t.LLMRotate, &t.FirstRunAt, &t.DeadlineAt, &t.TargetGate); err != nil {
return nil, err
}
t.Pinned = t.PinnedAt != nil
@@ -420,6 +426,8 @@ type TaskPatch struct {
Pinned *bool
// PlanQuietMode flips the per-task quiet-mode switch (#86); nil = leave unchanged.
PlanQuietMode *bool
// TargetGate flips the per-task target-gate opt-in; nil = leave unchanged.
TargetGate *bool
}
// UpdateTask applies a partial task name/pin mutation and returns the updated row.
@@ -432,9 +440,10 @@ func (d *DB) UpdateTask(id int64, patch TaskPatch) (*Task, error) {
WHEN $4::boolean THEN COALESCE(pinned_at, now())
ELSE NULL
END,
plan_quiet_mode = CASE WHEN $5::boolean THEN $6 ELSE plan_quiet_mode END
plan_quiet_mode = CASE WHEN $5::boolean THEN $6 ELSE plan_quiet_mode END,
target_gate = CASE WHEN $7::boolean THEN $8 ELSE target_gate END
WHERE id=$1 AND deleted_at IS NULL
RETURNING `+taskCols, id, patch.Name != nil, patch.Name, patch.Pinned, patch.PlanQuietMode != nil, patch.PlanQuietMode))
RETURNING `+taskCols, id, patch.Name != nil, patch.Name, patch.Pinned, patch.PlanQuietMode != nil, patch.PlanQuietMode, patch.TargetGate != nil, patch.TargetGate))
if err == sql.ErrNoRows {
return nil, nil
}
@@ -673,6 +682,14 @@ DELETE FROM assets a USING deletable d WHERE a.id=d.id`, id, expID)
}
// llm_usage (the token metering ledger) is intentionally NOT deleted with the
// task — it is kept as historical accounting even after the task is gone.
// Target pins and approval rows are kept as history too: the pins are revoked,
// the open target requests settled, in this transaction.
if err := revokeTaskPins(tx, id); err != nil {
return result, err
}
if err := settleTaskTargetRequests(tx, id); err != nil {
return result, err
}
if _, err := tx.Exec(`DELETE FROM tasks WHERE id=$1`, id); err != nil {
return result, err
}
+56 -12
View File
@@ -87,11 +87,14 @@ const DefaultTargetWaitBound = 300 * time.Second
// (TaskID) or one conversation (ConversationID), never both. Enabled reports the
// scope's opt-in flag and is read on every gated call; nil means the gate is off.
// WaitBound caps the wait for a human decision; 0 means never wait.
// Interceptor is the one that owns the gate of this session, so that code
// holding only a context (an HTTP redirect callback) can reach it.
type TargetScope struct {
TaskID string
ConversationID int64
Enabled func() bool
WaitBound time.Duration
Interceptor *Interceptor
}
// WithTargetScope returns a child context carrying the target-gate scope of the
@@ -112,36 +115,52 @@ type compiledRule struct {
re *regexp.Regexp
}
// pendingManager tracks in-flight "ask" requests via per-request channels.
// pendingManager tracks in-flight approval requests via per-waiter channels. A
// classic ask has one waiter; a target request can have several (every call
// that names the same host waits on the same row).
type pendingManager struct {
mu sync.Mutex
ch map[int64]chan bool
ch map[int64][]chan bool
}
func newPendingManager() *pendingManager { return &pendingManager{ch: map[int64]chan bool{}} }
func newPendingManager() *pendingManager { return &pendingManager{ch: map[int64][]chan bool{}} }
// add registers one more waiter for id and returns its channel.
func (p *pendingManager) add(id int64) chan bool {
ch := make(chan bool, 1)
p.mu.Lock()
p.ch[id] = ch
p.ch[id] = append(p.ch[id], ch)
p.mu.Unlock()
return ch
}
// resolve wakes every waiter of id with the decision.
func (p *pendingManager) resolve(id int64, allowed bool) {
p.mu.Lock()
ch, ok := p.ch[id]
waiters := p.ch[id]
delete(p.ch, id)
p.mu.Unlock()
if ok {
for _, ch := range waiters {
ch <- allowed
}
}
func (p *pendingManager) remove(id int64) {
// remove takes one waiter's channel out; the other waiters of id stay.
func (p *pendingManager) remove(id int64, ch chan bool) {
p.mu.Lock()
delete(p.ch, id)
p.mu.Unlock()
defer p.mu.Unlock()
waiters := p.ch[id]
for n, w := range waiters {
if w == ch {
waiters = append(waiters[:n:n], waiters[n+1:]...)
break
}
}
if len(waiters) == 0 {
delete(p.ch, id)
return
}
p.ch[id] = waiters
}
// Reviewer runs the LLM fallback judge for one tool call and returns its verdict
@@ -159,6 +178,18 @@ type Interceptor struct {
enabledTools map[string]bool // nil means not loaded yet
pending *pendingManager
reviewer Reviewer // nil = LLM fallback judge not wired
// upstreamProxy reports whether target traffic leaves through an upstream
// egress proxy (nil = no). Injected by the server layer.
upstreamProxy func() bool
}
// SetUpstreamProxyCheck installs the function that reports whether an upstream
// egress proxy is configured. Behind one the proxy chooses the address it
// connects to, so a target approval given while it is set is name-only.
func (i *Interceptor) SetUpstreamProxyCheck(f func() bool) {
i.mu.Lock()
i.upstreamProxy = f
i.mu.Unlock()
}
// SetReviewer installs the LLM fallback judge callback. Passing nil disables it.
@@ -633,7 +664,7 @@ func (i *Interceptor) HandleAsk(ctx context.Context, convID int64, dec Decision,
i.bindResult(ctx, pendingID, audit)
ch := i.pending.add(pendingID)
defer i.pending.remove(pendingID)
defer i.pending.remove(pendingID, ch)
// Polling clients can decide after INSERT commits but before the channel is
// registered. Re-read after registration so that decision cannot be lost;
// later decisions will be delivered through ch.
@@ -707,7 +738,11 @@ func (i *Interceptor) HandleAsk(ctx context.Context, convID int64, dec Decision,
var ErrAlreadyDecided = errors.New("approval already decided or does not exist; please refresh the list")
// Decide resolves a pending request. Called by the HTTP decide endpoint.
func (i *Interceptor) Decide(pendingID int64, allowed bool) error {
// approver (subject of the session token) and remoteAddr identify the human
// behind the decision; they are recorded in the row's audit. On a target
// request "allowed" also pins the addresses the request displayed, and both
// values go to the pin: a human decision is the only thing that creates a pin.
func (i *Interceptor) Decide(pendingID int64, allowed bool, approver, remoteAddr string) error {
status := "denied"
if allowed {
status = "allowed"
@@ -716,7 +751,16 @@ func (i *Interceptor) Decide(pendingID int64, allowed bool) error {
if allowed {
action, reason = "allow", "user approved execution"
}
resolved, err := i.db.ResolveIntercept(pendingID, status, action, reason, "human")
targetKey, err := i.db.TargetRequestKey(pendingID)
if err != nil {
return err
}
var resolved bool
if targetKey != "" {
_, resolved, err = i.db.ResolveTargetRequest(pendingID, allowed, approver, remoteAddr)
} else {
resolved, err = i.db.ResolveInterceptBy(pendingID, status, action, reason, "human", approver, remoteAddr)
}
if err != nil {
return err
}
+53
View File
@@ -0,0 +1,53 @@
package intercept
import (
"fmt"
"net/netip"
)
// metadataAddrs are instance-metadata addresses that lie outside the link-local
// ranges and would otherwise pass as private or public unicast.
var metadataAddrs = []netip.Addr{
netip.MustParseAddr("fd00:ec2::254"), // AWS IMDS over IPv6
netip.MustParseAddr("100.100.100.200"), // Alibaba Cloud
}
// ValidateExemptionHost checks the host of an entry an operator adds to the
// target-gate exemption list. host is expected in normalised form.
//
// An exemption on a name is resolved at check time and is void when the name
// resolves to a private, loopback, link-local or unspecified address. An IP
// literal is never resolved, so that rule cannot protect it and the literal is
// checked here instead. A private unicast address is allowed: an internal
// service the operator names on purpose (10.0.0.5:8080). Loopback, link-local
// (169.254.0.0/16 holds the cloud metadata address 169.254.169.254; fe80::/10),
// unspecified and multicast addresses are never accepted, nor are the metadata
// addresses outside the link-local ranges. A host that is not an IP literal
// passes: names are the void rule's business.
func ValidateExemptionHost(host string) error {
addr, err := netip.ParseAddr(host)
if err != nil {
return nil
}
addr = addr.Unmap().WithZone("")
kind := ""
switch {
case addr.IsLoopback():
kind = "a loopback"
case addr.IsUnspecified():
kind = "the unspecified"
case addr.IsLinkLocalUnicast():
kind = "a link-local (metadata range)"
case addr.IsMulticast():
kind = "a multicast"
}
for _, m := range metadataAddrs {
if addr == m {
kind = "an instance-metadata"
}
}
if kind != "" {
return fmt.Errorf("%s is %s address and can never be exempt from target approval", host, kind)
}
return nil
}
+29
View File
@@ -0,0 +1,29 @@
package intercept
import "testing"
// An operator may exempt a private unicast IP literal, never a loopback,
// link-local (cloud metadata), unspecified or multicast one.
func TestValidateExemptionHost(t *testing.T) {
for _, host := range []string{
"10.0.0.5", "172.16.0.1", "192.168.1.1", "fd00::1", // private unicast
"203.0.113.7", "2001:db8::1", // public
"pypi.org", "internal.corp", "localhost", // names: the void rule's business
} {
if err := ValidateExemptionHost(host); err != nil {
t.Errorf("%s refused: %v", host, err)
}
}
for _, host := range []string{
"169.254.169.254", "169.254.170.2", "169.254.0.1", "fe80::1", "fe80::1%eth0", // link-local, metadata
"::ffff:169.254.169.254", // the same address, IPv4-mapped
"fd00:ec2::254", "100.100.100.200", // metadata outside link-local
"127.0.0.1", "127.8.8.8", "::1", "::ffff:127.0.0.1", // loopback
"0.0.0.0", "::", // unspecified
"224.0.0.1", "239.255.255.250", "ff02::1", "ff05::2", // multicast
} {
if err := ValidateExemptionHost(host); err == nil {
t.Errorf("%s accepted, want it refused", host)
}
}
}
+374
View File
@@ -0,0 +1,374 @@
package intercept
import (
"context"
"encoding/json"
"fmt"
"net"
"strconv"
"strings"
"time"
"github.com/Autumn-27/artex/db"
)
// Block codes a target request can end with. They are carried in the
// model-facing text of the blocked call.
const (
// TargetApprovalPending: no decision within the wait bound. The request
// stays open, so a later approval still creates the pin.
TargetApprovalPending = "TARGET_APPROVAL_PENDING"
// TargetDenied: the request was closed without an approval.
TargetDenied = "TARGET_DENIED"
// TargetGateUnavailable: the request could not be written or read.
TargetGateUnavailable = "TARGET_GATE_UNAVAILABLE"
)
// targetPollInterval is how often a waiter re-reads its request row. The poll
// covers a decision taken on another instance that shares the database.
var targetPollInterval = 2 * time.Second
// TargetRequest is one target of a gated call that no pin covers.
type TargetRequest struct {
Host string // normalised name, IP literal or CIDR; with Port the key of the request
Port int // 0 = port not declared
Kind string
// ResolvedIPs are the addresses resolved for this request. They are what the
// human reads and what an approval pins.
ResolvedIPs []string
Resolver string
ResolvedAt time.Time
Reason string // no_pin | new_ip | revoked
Pin *db.TargetPin // the active pin an approval would extend (new_ip); nil otherwise
}
// TargetResult is the outcome of RequestTargets. Code is empty when every
// request the call waited on was approved (or was decided in between): the
// caller then re-runs the pin lookup and executes only if every target is
// covered. Otherwise Code is the block code and Host, Port and PendingID name
// the request that stopped the call.
type TargetResult struct {
Code string
Host string
Port int
PendingID int64
Reason string
// AskMerged reports that the ask passed in was carried by a request this call
// created, so its decision answered both. When false the caller still has to
// raise the ask on its own.
AskMerged bool
}
func (s TargetScope) ref() (db.TargetScopeRef, error) {
ref := db.TargetScopeRef{ConversationID: s.ConversationID}
if s.TaskID != "" {
id, err := strconv.ParseInt(s.TaskID, 10, 64)
if err != nil {
return ref, fmt.Errorf("target scope: task id %q: %w", s.TaskID, err)
}
ref.TaskID = id
}
return ref, nil
}
// scopeName is how a request text names what an approval is bound to.
func (s TargetScope) scopeName() string {
if s.TaskID != "" {
return "task " + s.TaskID
}
return "conversation " + strconv.FormatInt(s.ConversationID, 10)
}
// targetWaiter is one request row a call waits on.
type targetWaiter struct {
target TargetRequest
id int64 // 0 = the open request was decided in between: nothing to wait for
created bool // this call wrote the row
}
// RequestTargets asks a human to approve the targets of one gated call that no
// pin covers, and waits for the decisions up to the scope's wait bound.
//
// There is one open request per (scope, host, port). The first call writes the
// row and emits the approval card once; later calls for the same key attach to
// that row and wait on it. A request is durable and host-level: a call that
// stops waiting (bound elapsed, cancelled) leaves it open, and a later approval
// still creates the pin. Nothing here ever pins: only Decide does, on a click.
//
// ask is the pending ask of a rule or of the judge for the same call, nil when
// there is none. It is merged into the first row this call creates, so one
// click answers both. A merged row is bound to its creating call like a classic
// ask: when that call leaves, the system settles it, and a rule timeout never
// allows it, because a timeout cannot pin.
func (i *Interceptor) RequestTargets(ctx context.Context, scope TargetScope, toolName string, input []byte, targets []TargetRequest, ask *Decision) TargetResult {
ref, err := scope.ref()
if err != nil {
return TargetResult{Code: TargetGateUnavailable, Reason: err.Error()}
}
if scope.WaitBound <= 0 {
// A call that never waits cannot own a call-bound row.
ask = nil
}
var res TargetResult
var merged *targetWaiter
var mergedAsk Decision
// Whatever the way out, a merged row never outlives its creating call. This
// is a no-op once the row is decided.
defer func() {
if merged != nil {
i.settleMerged(merged.id, "denied", "the call left before a decision")
}
}()
waiters := make([]*targetWaiter, 0, len(targets))
for _, t := range targets {
w, err := i.openTargetRequest(ctx, ref, scope, toolName, input, t, ask)
if err != nil {
res.Code, res.Host, res.Port, res.Reason = TargetGateUnavailable, t.Host, t.Port, err.Error()
return res
}
if w.created && ask != nil {
merged, mergedAsk, res.AskMerged = w, *ask, true
ask = nil
}
waiters = append(waiters, w)
}
if len(waiters) == 0 {
return res
}
blocked := func(w *targetWaiter, code, reason string) TargetResult {
res.Code, res.Host, res.Port, res.PendingID, res.Reason = code, w.target.Host, w.target.Port, w.id, reason
return res
}
if scope.WaitBound <= 0 {
return blocked(waiters[0], TargetApprovalPending, "this call does not wait for approvals; the request stays open")
}
bound := time.NewTimer(scope.WaitBound)
defer bound.Stop()
// The merged row first: its rule timeout runs from the creation of the row.
if merged != nil {
var ruleTimeout <-chan time.Time
if mergedAsk.TimeoutEnabled {
secs := mergedAsk.TimeoutSeconds
if secs <= 0 {
secs = 60
}
timer := time.NewTimer(time.Duration(secs) * time.Second)
defer timer.Stop()
ruleTimeout = timer.C
}
status, reason := "", ""
switch st := i.awaitTarget(ctx, merged.id, bound.C, ruleTimeout); st {
case "allowed":
case waitRuleTimeout:
status, reason = "timeout", "approval timed out"
if mergedAsk.TimeoutAction == "allow" {
reason += `; the timeout action "allow" was overridden, because a timeout cannot pin a target`
}
case waitElapsed:
status, reason = "timeout", "no approval within the wait bound"
case waitCancelled:
status, reason = "denied", "work cancelled"
case waitGone:
return blocked(merged, TargetGateUnavailable, "the approval request no longer exists")
default:
return blocked(merged, TargetDenied, "the request was denied")
}
// A human decision that lands while the call is leaving still counts.
if status != "" && (i.settleMerged(merged.id, status, reason) || !i.targetAllowed(merged.id)) {
return blocked(merged, TargetDenied, reason)
}
}
for _, w := range waiters {
if w == merged || w.id == 0 {
continue
}
st := i.awaitTarget(ctx, w.id, bound.C, nil)
if (st == "denied" || st == "timeout") && !w.created && i.wasMerged(w.id) {
// The row this call attached to was bound to another call and closed
// with it. That says nothing about the host: ask once more, with a
// plain request.
retry, err := i.openTargetRequest(ctx, ref, scope, toolName, input, w.target, nil)
if err != nil {
return blocked(w, TargetGateUnavailable, err.Error())
}
if retry.id == 0 {
continue
}
w = retry
st = i.awaitTarget(ctx, w.id, bound.C, nil)
}
switch st {
case "allowed":
case waitElapsed:
return blocked(w, TargetApprovalPending, "no approval within the wait bound; the request stays open")
case waitCancelled:
return blocked(w, TargetApprovalPending, "work cancelled; the request stays open")
case waitGone:
return blocked(w, TargetGateUnavailable, "the approval request no longer exists")
default:
return blocked(w, TargetDenied, "the request was denied")
}
}
return res
}
// openTargetRequest finds or creates the request row for one target. When this
// call writes the row it emits the approval card; an attaching call writes and
// emits nothing, and never rewrites the addresses a human may be reading.
func (i *Interceptor) openTargetRequest(ctx context.Context, ref db.TargetScopeRef, scope TargetScope, toolName string, input []byte, t TargetRequest, ask *Decision) (*targetWaiter, error) {
i.mu.RLock()
proxied := i.upstreamProxy
i.mu.RUnlock()
// Behind the upstream proxy the proxy chooses the address: the approval is
// for the name, and the pin created from this row says so.
nameOnly := proxied != nil && proxied()
dec := Decision{Action: "ask"}
if ask != nil {
dec = *ask
}
dec.Message = targetReason(scope, t, ask, nameOnly)
audit := auditFor(ctx, dec, input, "pending")
audit.MergedAsk = ask != nil
target := db.InterceptTarget{
Host: t.Host, Port: t.Port, Kind: t.Kind,
ResolvedIPs: t.ResolvedIPs, Resolver: t.Resolver, ResolvedAt: t.ResolvedAt,
Reason: t.Reason, Outcome: "requested", NameOnlyEffective: nameOnly,
}
if t.Pin != nil {
target.PinID, target.PinVersion, target.PinnedIPs = t.Pin.ID, t.Pin.Version, t.Pin.PinnedIPs
}
audit.Targets = []db.InterceptTarget{target}
_, agentName := taskInfoFromCtx(ctx)
id, created, err := i.db.FindOrCreateTargetRequest(db.TargetRequest{
Scope: ref, Host: t.Host, Port: t.Port, RuleID: dec.RuleID,
AgentName: agentName, ToolName: toolName, ToolInput: input, Reason: dec.Message, Audit: audit,
})
if err != nil {
return nil, err
}
w := &targetWaiter{target: t, id: id, created: created}
if !created {
return w, nil
}
if ask != nil {
// Call-bound like a classic ask: the result of this call belongs to the row.
i.bindResult(ctx, id, audit)
}
detail, _ := json.Marshal(map[string]any{
"pending_id": id,
"tool": toolName,
"input": json.RawMessage(input),
})
activity := db.Activity{
Kind: "intercept_request",
Summary: fmt.Sprintf("tool %s requests approval for target %s (#%d)", toolName, targetName(t), id),
Detail: string(detail),
}
if scope.ConversationID != 0 {
_, _ = i.db.AppendConvActivity(scope.ConversationID, activity)
} else if emit := taskEmitFromCtx(ctx); emit != nil {
emit(activity)
}
return w, nil
}
// Ways a wait ends without a decision on the row.
const (
waitElapsed = "wait_elapsed" // the scope's wait bound
waitRuleTimeout = "wait_rule_timeout" // the timeout of a merged ask
waitCancelled = "wait_cancelled" // the call's context
waitGone = "wait_gone" // the row no longer exists (its conversation was deleted)
)
// awaitTarget waits for the decision on one request row and returns its status
// (allowed | denied | timeout) or the way the wait ended. It never settles the
// row. The row is re-read before each wait: that covers a decision taken before
// the waiter registered, and one taken on another instance.
func (i *Interceptor) awaitTarget(ctx context.Context, id int64, bound, ruleTimeout <-chan time.Time) string {
ch := i.pending.add(id)
defer i.pending.remove(id, ch)
poll := time.NewTicker(targetPollInterval)
defer poll.Stop()
for {
row, err := i.db.GetInterceptPending(id)
if err == nil && row == nil {
return waitGone
}
if err == nil && row.Status != "pending" {
return row.Status
}
select {
case <-ch:
case <-poll.C:
case <-bound:
return waitElapsed
case <-ruleTimeout:
return waitRuleTimeout
case <-ctx.Done():
return waitCancelled
}
}
}
// settleMerged closes a merged row on behalf of its creating call, which is
// leaving, and wakes the calls attached to it. It reports whether the row was
// still open; a merged row closed this way never pins.
func (i *Interceptor) settleMerged(id int64, status, reason string) bool {
resolved, err := i.db.ResolveIntercept(id, status, "deny", reason, "system")
if err != nil || !resolved {
return false
}
i.pending.resolve(id, false)
return true
}
func (i *Interceptor) targetAllowed(id int64) bool {
row, err := i.db.GetInterceptPending(id)
return err == nil && row != nil && row.Status == "allowed"
}
func (i *Interceptor) wasMerged(id int64) bool {
detail, err := i.db.GetInterceptDetail(id)
return err == nil && detail != nil && detail.Audit != nil && detail.Audit.MergedAsk
}
func targetName(t TargetRequest) string {
if t.Port == 0 {
return t.Host
}
return net.JoinHostPort(t.Host, strconv.Itoa(t.Port))
}
// targetReason is the text a human reads before deciding: host, port, the
// addresses an approval pins, and what the approval is bound to.
func targetReason(scope TargetScope, t TargetRequest, ask *Decision, nameOnly bool) string {
var b strings.Builder
if ask != nil && ask.Message != "" {
b.WriteString(ask.Message)
b.WriteString("\nThis call also names a target that needs approval. ")
}
name, bound := targetName(t), scope.scopeName()
if t.Port == 0 {
name += " (port not declared: every port)"
}
ips := strings.Join(t.ResolvedIPs, ", ")
if t.Pin != nil {
fmt.Fprintf(&b, "Target %s is pinned for %s (pin v%d: %s) and now resolves to %s (resolver: %s).",
name, bound, t.Pin.Version, strings.Join(t.Pin.PinnedIPs, ", "), ips, t.Resolver)
} else {
fmt.Fprintf(&b, "Target %s is not approved for %s. It resolves to %s (resolver: %s).", name, bound, ips, t.Resolver)
}
fmt.Fprintf(&b, " Approving pins these addresses for %s until revoked.", bound)
if nameOnly {
b.WriteString(" Egress goes through the upstream proxy, which chooses the address it connects to." +
" The addresses listed are what the platform resolved. This approval is name-only.")
}
return b.String()
}
+754
View File
@@ -0,0 +1,754 @@
package intercept
import (
"context"
"fmt"
"slices"
"strconv"
"strings"
"sync/atomic"
"testing"
"time"
"github.com/Autumn-27/artex/db"
)
func TestPendingWakesAllWaiters(t *testing.T) {
p := newPendingManager()
a, b, gone := p.add(1), p.add(1), p.add(1)
other := p.add(2)
p.remove(1, gone)
p.resolve(1, true)
for name, ch := range map[string]chan bool{"first": a, "second": b} {
select {
case allowed := <-ch:
if !allowed {
t.Fatalf("%s waiter got a denial", name)
}
default:
t.Fatalf("%s waiter of the request was not woken", name)
}
}
select {
case <-gone:
t.Fatal("a removed waiter was woken")
case <-other:
t.Fatal("a waiter of another request was woken")
default:
}
if len(p.ch) != 1 || len(p.ch[2]) != 1 {
t.Fatalf("waiters left = %v, want only the one of request 2", p.ch)
}
// Resolving a request nobody waits on, twice, blocks nobody.
p.resolve(1, true)
p.resolve(2, false)
if allowed := <-other; allowed {
t.Fatal("the waiter of request 2 got an approval")
}
p.remove(2, other)
if len(p.ch) != 0 {
t.Fatalf("waiters left = %v, want none", p.ch)
}
}
func targetTestDB(t *testing.T) *db.DB {
t.Helper()
dsn, _, err := db.DSN()
if err != nil {
t.Skipf("no database config (%v) - skipping", err)
}
d, err := db.Open(dsn)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { d.Close() })
return d
}
// taskScope creates a task and returns its target scope on ic, with the gate
// on and the given wait bound.
func taskScope(t *testing.T, d *db.DB, ic *Interceptor, wait time.Duration) TargetScope {
t.Helper()
task, err := d.CreateTaskWithOptions("target requests", "request, wait, decide", db.TaskCreateOptions{})
if err != nil {
t.Fatal(err)
}
id := strconv.FormatInt(task.ID, 10)
t.Cleanup(func() {
_ = d.DeleteTask(task.ID)
_, _ = d.Exec(`DELETE FROM target_pins WHERE task_id=$1`, task.ID)
_, _ = d.Exec(`DELETE FROM intercept_pending WHERE task_id=$1`, id)
})
return TargetScope{TaskID: id, Enabled: func() bool { return true }, WaitBound: wait, Interceptor: ic}
}
func (s TargetScope) dbRef(t *testing.T) db.TargetScopeRef {
t.Helper()
ref, err := s.ref()
if err != nil {
t.Fatal(err)
}
return ref
}
func target(host string, port int, ips ...string) TargetRequest {
return TargetRequest{Host: host, Port: port, Kind: "domain", ResolvedIPs: ips, Resolver: "system", ResolvedAt: time.Now().UTC(), Reason: "no_pin"}
}
var webFetchInput = []byte(`{"url":"https://app.example.com/"}`)
// request runs RequestTargets in the background and returns the channel its
// result arrives on.
func request(ctx context.Context, ic *Interceptor, scope TargetScope, ask *Decision, targets ...TargetRequest) <-chan TargetResult {
out := make(chan TargetResult, 1)
go func() { out <- ic.RequestTargets(ctx, scope, "WebFetch", webFetchInput, targets, ask) }()
return out
}
func result(t *testing.T, ch <-chan TargetResult) TargetResult {
t.Helper()
select {
case res := <-ch:
return res
case <-time.After(20 * time.Second):
t.Fatal("the call is still waiting")
return TargetResult{}
}
}
func stillWaiting(t *testing.T, ch <-chan TargetResult) {
t.Helper()
select {
case res := <-ch:
t.Fatalf("the call returned %+v, want it still waiting", res)
case <-time.After(4 * targetPollInterval):
}
}
func eventually(t *testing.T, what string, cond func() bool) {
t.Helper()
deadline := time.Now().Add(20 * time.Second)
for !cond() {
if time.Now().After(deadline) {
t.Fatalf("timed out waiting for %s", what)
}
time.Sleep(5 * time.Millisecond)
}
}
// openRows returns the ids of the open target requests of scope, oldest first.
func openRows(t *testing.T, d *db.DB, scope TargetScope) []int64 {
t.Helper()
rows, err := d.Query(`SELECT id FROM intercept_pending
WHERE COALESCE(task_id,'')=$1 AND COALESCE(conversation_id,0)=$2 AND status='pending' AND target_key IS NOT NULL ORDER BY id`,
scope.TaskID, scope.ConversationID)
if err != nil {
t.Fatal(err)
}
defer rows.Close()
var ids []int64
for rows.Next() {
var id int64
if err := rows.Scan(&id); err != nil {
t.Fatal(err)
}
ids = append(ids, id)
}
return ids
}
func allRows(t *testing.T, d *db.DB, scope TargetScope) int {
t.Helper()
var n int
if err := d.QueryRow(`SELECT count(*) FROM intercept_pending
WHERE COALESCE(task_id,'')=$1 AND COALESCE(conversation_id,0)=$2 AND target_key IS NOT NULL`, scope.TaskID, scope.ConversationID).Scan(&n); err != nil {
t.Fatal(err)
}
return n
}
// waiters is how many calls wait on request id through ic.
func waiters(ic *Interceptor, id int64) int {
ic.pending.mu.Lock()
defer ic.pending.mu.Unlock()
return len(ic.pending.ch[id])
}
// oneOpenRow waits for the scope's single open request and returns its id.
func oneOpenRow(t *testing.T, d *db.DB, scope TargetScope) int64 {
t.Helper()
var ids []int64
eventually(t, "one open target request", func() bool {
ids = openRows(t, d, scope)
return len(ids) == 1
})
return ids[0]
}
func pins(t *testing.T, d *db.DB, scope TargetScope) []*db.TargetPin {
t.Helper()
out, err := d.ListTargetPins(scope.dbRef(t))
if err != nil {
t.Fatal(err)
}
return out
}
// Calls for one host share one request, in one process and across two
// instances on one database: one row, one card, one decision wakes them all,
// one pin.
func TestConcurrentCallsSameHostShareOneRequest(t *testing.T) {
d := targetTestDB(t)
first, second := New(d), New(targetTestDB(t))
scope := taskScope(t, d, first, time.Minute)
var cards atomic.Int32
ctx := WithTaskContext(context.Background(), scope.TaskID, "work#1", func(a db.Activity) {
if a.Kind == "intercept_request" {
cards.Add(1)
}
})
tgt := target("app.example.com", 443, "203.0.113.7")
a := request(ctx, first, scope, nil, tgt)
id := oneOpenRow(t, d, scope)
b := request(ctx, first, scope, nil, tgt)
c := request(ctx, second, scope, nil, tgt)
eventually(t, "every call to wait on the request", func() bool { return waiters(first, id) == 2 && waiters(second, id) == 1 })
if ids := openRows(t, d, scope); len(ids) != 1 || allRows(t, d, scope) != 1 {
t.Fatalf("open requests = %v, want the one row %d", ids, id)
}
if n := cards.Load(); n != 1 {
t.Fatalf("%d approval cards emitted, want 1", n)
}
stillWaiting(t, a)
// Decided on the second instance: its own waiter is woken by the channel,
// the waiters of the first instance by their poll.
if err := second.Decide(id, true, "ARTEX", "198.51.100.9"); err != nil {
t.Fatal(err)
}
for name, ch := range map[string]<-chan TargetResult{"creator": a, "attached": b, "other instance": c} {
if res := result(t, ch); res.Code != "" {
t.Fatalf("%s call: %+v, want approved", name, res)
}
}
got := pins(t, d, scope)
if len(got) != 1 || got[0].Status != db.TargetPinActive || got[0].Host != "app.example.com" || got[0].Port != 443 ||
!slices.Equal(got[0].PinnedIPs, []string{"203.0.113.7"}) {
t.Fatalf("pins = %+v, want one active pin", got)
}
if allRows(t, d, scope) != 1 || cards.Load() != 1 {
t.Fatalf("rows=%d cards=%d after the decision, want 1 and 1", allRows(t, d, scope), cards.Load())
}
if err := first.Decide(id, false, "ARTEX", ""); err != ErrAlreadyDecided {
t.Fatalf("second decision: %v, want ErrAlreadyDecided", err)
}
}
// A call that stops waiting leaves the request alone: it stays open, the next
// call attaches to it, and a later approval still pins.
func TestWaiterLeavingKeepsRequestOpen(t *testing.T) {
d := targetTestDB(t)
ic := New(d)
scope := taskScope(t, d, ic, time.Minute)
tgt := target("app.example.com", 443, "203.0.113.7")
ctx, cancel := context.WithCancel(context.Background())
call := request(ctx, ic, scope, nil, tgt)
id := oneOpenRow(t, d, scope)
eventually(t, "the call to wait", func() bool { return waiters(ic, id) == 1 })
cancel()
res := result(t, call)
if res.Code != TargetApprovalPending || res.PendingID != id || res.Host != "app.example.com" || res.Port != 443 {
t.Fatalf("cancelled call: %+v, want %s on request %d", res, TargetApprovalPending, id)
}
row, err := d.GetInterceptPending(id)
if err != nil || row == nil || row.Status != "pending" || row.DecidedBy != "" {
t.Fatalf("request after its waiter left = %+v, %v; want untouched", row, err)
}
if waiters(ic, id) != 0 {
t.Fatal("the waiter that left is still registered")
}
again := request(context.Background(), ic, scope, nil, tgt)
eventually(t, "the next call to attach", func() bool { return waiters(ic, id) == 1 })
if ids := openRows(t, d, scope); !slices.Equal(ids, []int64{id}) {
t.Fatalf("open requests = %v, want the same row %d", ids, id)
}
if err := ic.Decide(id, true, "ARTEX", ""); err != nil {
t.Fatal(err)
}
if res := result(t, again); res.Code != "" {
t.Fatalf("attached call: %+v, want approved", res)
}
if got := pins(t, d, scope); len(got) != 1 {
t.Fatalf("pins = %+v, want the late approval to pin", got)
}
}
func TestWaitBoundReturnsPendingCode(t *testing.T) {
d := targetTestDB(t)
ic := New(d)
scope := taskScope(t, d, ic, 150*time.Millisecond)
tgt := target("app.example.com", 443, "203.0.113.7")
start := time.Now()
res := result(t, request(context.Background(), ic, scope, nil, tgt))
if res.Code != TargetApprovalPending || res.PendingID == 0 {
t.Fatalf("bounded call: %+v, want %s", res, TargetApprovalPending)
}
if waited := time.Since(start); waited < 150*time.Millisecond {
t.Fatalf("the call returned after %v, before its wait bound", waited)
}
if ids := openRows(t, d, scope); !slices.Equal(ids, []int64{res.PendingID}) {
t.Fatalf("open requests = %v, want the request %d left open", ids, res.PendingID)
}
// Bound 0 (the final settlement round): never wait, but still ask.
never := scope
never.WaitBound = 0
other := target("other.example.com", 443, "203.0.113.8")
res = result(t, request(context.Background(), ic, never, nil, other))
if res.Code != TargetApprovalPending || res.Host != "other.example.com" || res.PendingID == 0 {
t.Fatalf("zero-bound call: %+v, want an immediate %s", res, TargetApprovalPending)
}
if len(openRows(t, d, scope)) != 2 {
t.Fatalf("open requests = %v, want the zero-bound call to have created its request", openRows(t, d, scope))
}
// An ask is never merged into a row the call does not wait on.
asked := ic.RequestTargets(context.Background(), never, "WebFetch", webFetchInput,
[]TargetRequest{target("third.example.com", 443, "203.0.113.9")}, &Decision{Action: "ask", Message: "rule asks"})
if asked.Code != TargetApprovalPending || asked.AskMerged {
t.Fatalf("zero-bound call with an ask: %+v, want a plain open request", asked)
}
if row, err := d.GetInterceptPending(asked.PendingID); err != nil || row == nil || row.Status != "pending" {
t.Fatalf("request of the zero-bound call = %+v, %v; want left open", row, err)
}
}
func askRule(t *testing.T, d *db.DB, timeoutEnabled bool, seconds int, action string) Decision {
t.Helper()
rule, err := d.CreateInterceptRule(fmt.Sprintf("target merge %d", time.Now().UnixNano()), "tool_name", "string",
"NoSuchToolForTargetMerge", "ask", "", 0, false, timeoutEnabled, seconds, action)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = d.DeleteInterceptRule(rule.ID) })
return Decision{Action: "ask", Message: "rule requires approval", RuleID: rule.ID, RuleName: rule.Name,
TimeoutEnabled: timeoutEnabled, TimeoutSeconds: seconds, TimeoutAction: action}
}
// A merged row is bound to the call that created it: when that call leaves, the
// system closes the row, and nothing is pinned.
func TestMergedRequestSettledWhenCreatorLeaves(t *testing.T) {
d := targetTestDB(t)
ic := New(d)
scope := taskScope(t, d, ic, time.Minute)
ask := askRule(t, d, false, 0, "deny")
ctx, cancel := context.WithCancel(context.Background())
call := request(ctx, ic, scope, &ask, target("app.example.com", 443, "203.0.113.7"))
id := oneOpenRow(t, d, scope)
detail, err := d.GetInterceptDetail(id)
if err != nil || detail == nil || detail.RuleID == nil || *detail.RuleID != ask.RuleID || detail.DecisionSource != "target" ||
detail.Audit == nil || !detail.Audit.MergedAsk || detail.Audit.RuleName != ask.RuleName || len(detail.Audit.Targets) != 1 {
t.Fatalf("merged row = %+v, %v; want the rule and the target on one row", detail, err)
}
if !strings.Contains(detail.Reason, "rule requires approval") || !strings.Contains(detail.Reason, "app.example.com:443") ||
!strings.Contains(detail.Reason, "203.0.113.7") || !strings.Contains(detail.Reason, "Approving pins these addresses for task "+scope.TaskID+" until revoked.") {
t.Fatalf("merged reason does not name both the ask and the target: %q", detail.Reason)
}
eventually(t, "the creating call to wait", func() bool { return waiters(ic, id) == 1 })
cancel()
res := result(t, call)
if res.Code != TargetDenied || !res.AskMerged || res.PendingID != id {
t.Fatalf("creating call after cancel: %+v, want %s with the ask merged", res, TargetDenied)
}
row, err := d.GetInterceptDetail(id)
if err != nil || row == nil || row.Status != "denied" || row.DecidedBy != "system" || row.Audit.EffectiveAction != "deny" ||
row.Audit.DecisionReason != "work cancelled" {
t.Fatalf("merged row after its creator left = %+v, %v; want denied by the system", row, err)
}
if got := pins(t, d, scope); len(got) != 0 {
t.Fatalf("a row closed by the system pinned: %+v", got)
}
}
// A timeout never pins, whatever the rule says: timeout_action = allow is
// overridden on a merged row and the override is recorded.
func TestTimeoutNeverPins(t *testing.T) {
d := targetTestDB(t)
ic := New(d)
scope := taskScope(t, d, ic, time.Minute)
ask := askRule(t, d, true, 1, "allow")
res := result(t, request(context.Background(), ic, scope, &ask, target("app.example.com", 443, "203.0.113.7")))
if res.Code != TargetDenied || !res.AskMerged {
t.Fatalf("call after the rule timeout: %+v, want %s", res, TargetDenied)
}
row, err := d.GetInterceptDetail(res.PendingID)
if err != nil || row == nil || row.Status != "timeout" || row.DecidedBy != "system" || row.Audit.EffectiveAction != "deny" ||
row.Audit.ExecutionStatus != "not_executed" || !strings.Contains(row.Audit.DecisionReason, `timeout action "allow" was overridden`) {
t.Fatalf("timed-out merged row = %+v, %v; want a timeout that denies and records the override", row, err)
}
if got := pins(t, d, scope); len(got) != 0 {
t.Fatalf("a timeout pinned: %+v", got)
}
if len(openRows(t, d, scope)) != 0 {
t.Fatal("the merged row is still open after its creator left")
}
}
// What is displayed is what is pinned: the addresses stored when the row was
// created. A call that attaches with other addresses never rewrites them.
func TestDecidePinsDisplayedIPs(t *testing.T) {
d := targetTestDB(t)
ic := New(d)
scope := taskScope(t, d, ic, time.Minute)
creator := request(context.Background(), ic, scope, nil, target("app.example.com", 443, "203.0.113.8", "203.0.113.7"))
id := oneOpenRow(t, d, scope)
attached := request(context.Background(), ic, scope, nil, target("app.example.com", 443, "198.51.100.66"))
eventually(t, "both calls to wait", func() bool { return waiters(ic, id) == 2 })
shown, err := d.GetInterceptDetail(id)
if err != nil || shown == nil || len(shown.Audit.Targets) != 1 ||
!slices.Equal(shown.Audit.Targets[0].ResolvedIPs, []string{"203.0.113.8", "203.0.113.7"}) || strings.Contains(shown.Reason, "198.51.100.66") {
t.Fatalf("displayed request = %+v, %v; want the creator's addresses only", shown, err)
}
if err := ic.Decide(id, true, "ARTEX", "198.51.100.9"); err != nil {
t.Fatal(err)
}
if res := result(t, creator); res.Code != "" {
t.Fatalf("creator: %+v", res)
}
if res := result(t, attached); res.Code != "" {
t.Fatalf("attached: %+v", res)
}
got := pins(t, d, scope)
if len(got) != 1 || !slices.Equal(got[0].PinnedIPs, []string{"203.0.113.7", "203.0.113.8"}) || got[0].InterceptID == nil || *got[0].InterceptID != id {
t.Fatalf("pins = %+v, want exactly the displayed addresses", got)
}
}
func TestDeniedRequestPinsNothing(t *testing.T) {
d := targetTestDB(t)
ic := New(d)
scope := taskScope(t, d, ic, time.Minute)
tgt := target("app.example.com", 443, "203.0.113.7")
call := request(context.Background(), ic, scope, nil, tgt)
id := oneOpenRow(t, d, scope)
if err := ic.Decide(id, false, "ARTEX", "198.51.100.9"); err != nil {
t.Fatal(err)
}
if res := result(t, call); res.Code != TargetDenied || res.PendingID != id {
t.Fatalf("denied call: %+v, want %s", res, TargetDenied)
}
if got := pins(t, d, scope); len(got) != 0 {
t.Fatalf("a denial pinned: %+v", got)
}
row, err := d.GetInterceptDetail(id)
if err != nil || row == nil || row.Status != "denied" || row.DecidedBy != "human" || row.Audit.Targets[0].Outcome != "denied" ||
row.Audit.Approver != "ARTEX" || row.Audit.ApproverAddr != "198.51.100.9" {
t.Fatalf("denied row = %+v, %v", row, err)
}
// No sticky denial: the next call asks again, on a new row.
next := request(context.Background(), ic, scope, nil, tgt)
if fresh := oneOpenRow(t, d, scope); fresh == id {
t.Fatal("the next call reused the denied row")
} else if err := ic.Decide(fresh, true, "ARTEX", ""); err != nil {
t.Fatal(err)
}
if res := result(t, next); res.Code != "" {
t.Fatalf("call after the denial: %+v, want approved", res)
}
}
// A call attached to a merged row that closes without an approval is not
// blocked for the host: the row was about another call. It asks once more with
// a plain request, then blocks.
func TestAttachedWaiterRequeuesAfterMergedDenial(t *testing.T) {
d := targetTestDB(t)
for _, approveRetry := range []bool{true, false} {
t.Run(fmt.Sprintf("retry approved=%v", approveRetry), func(t *testing.T) {
ic := New(d)
scope := taskScope(t, d, ic, time.Minute)
ask := askRule(t, d, false, 0, "deny")
tgt := target("app.example.com", 443, "203.0.113.7")
creator := request(context.Background(), ic, scope, &ask, tgt)
merged := oneOpenRow(t, d, scope)
attached := request(context.Background(), ic, scope, nil, tgt)
eventually(t, "both calls to wait", func() bool { return waiters(ic, merged) == 2 })
if err := ic.Decide(merged, false, "ARTEX", ""); err != nil {
t.Fatal(err)
}
if res := result(t, creator); res.Code != TargetDenied || !res.AskMerged {
t.Fatalf("creator of the merged row: %+v, want %s", res, TargetDenied)
}
var retry int64
eventually(t, "the attached call to ask again", func() bool {
ids := openRows(t, d, scope)
if len(ids) == 1 && ids[0] != merged {
retry = ids[0]
}
return retry != 0
})
plain, err := d.GetInterceptDetail(retry)
if err != nil || plain == nil || plain.RuleID != nil || plain.Audit.MergedAsk {
t.Fatalf("retry row = %+v, %v; want a plain target request", plain, err)
}
stillWaiting(t, attached)
if err := ic.Decide(retry, approveRetry, "ARTEX", ""); err != nil {
t.Fatal(err)
}
res := result(t, attached)
if approveRetry && (res.Code != "" || len(pins(t, d, scope)) != 1) {
t.Fatalf("attached call after its retry was approved: %+v pins=%+v", res, pins(t, d, scope))
}
if !approveRetry && (res.Code != TargetDenied || res.PendingID != retry || len(pins(t, d, scope)) != 0) {
t.Fatalf("attached call after its retry was denied: %+v, want %s", res, TargetDenied)
}
// One retry, then block: never a third request.
if n := allRows(t, d, scope); n != 2 {
t.Fatalf("%d request rows, want the merged one and one retry", n)
}
})
}
}
// The request rows of a conversation go with it (cascade). A waiter that finds
// its row gone stops with a block.
func TestWaiterStopsWhenRowDeleted(t *testing.T) {
d := targetTestDB(t)
ic := New(d)
conv, err := d.CreateConversation("assistant", fmt.Sprintf("target wait %d", time.Now().UnixNano()), nil)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
_ = d.DeleteConversation(conv.ID)
_, _ = d.Exec(`DELETE FROM target_pins WHERE conversation_id=$1`, conv.ID)
})
scope := TargetScope{ConversationID: conv.ID, Enabled: func() bool { return true }, WaitBound: time.Minute, Interceptor: ic}
call := request(WithConvID(context.Background(), conv.ID), ic, scope, nil, target("app.example.com", 443, "203.0.113.7"))
id := oneOpenRow(t, d, scope)
eventually(t, "the call to wait", func() bool { return waiters(ic, id) == 1 })
// The card went to the conversation stream.
var cards int
if err := d.QueryRow(`SELECT count(*) FROM conversation_activities WHERE conversation_id=$1 AND kind='intercept_request'`, conv.ID).Scan(&cards); err != nil || cards != 1 {
t.Fatalf("%d approval cards in the conversation (%v), want 1", cards, err)
}
if err := d.DeleteConversation(conv.ID); err != nil {
t.Fatal(err)
}
if res := result(t, call); res.Code != TargetGateUnavailable || res.PendingID != id {
t.Fatalf("call whose request was deleted: %+v, want %s", res, TargetGateUnavailable)
}
}
// Every human decision records who took it and from where: on the pin and in
// the audit of a target row, in the audit of a classic ask. decided_by on the
// row stays the actor class.
func TestDecideRecordsApproverAndAddress(t *testing.T) {
d := targetTestDB(t)
ic := New(d)
scope := taskScope(t, d, ic, time.Minute)
call := request(context.Background(), ic, scope, nil, target("app.example.com", 443, "203.0.113.7"))
id := oneOpenRow(t, d, scope)
if err := ic.Decide(id, true, "ARTEX", "198.51.100.9"); err != nil {
t.Fatal(err)
}
if res := result(t, call); res.Code != "" {
t.Fatalf("approved call: %+v", res)
}
got := pins(t, d, scope)
if len(got) != 1 || got[0].DecidedBy != "ARTEX" || got[0].DecidedFrom != "198.51.100.9" {
t.Fatalf("pin = %+v, want the approver and the address", got)
}
row, err := d.GetInterceptDetail(id)
if err != nil || row == nil || row.DecidedBy != "human" || row.Audit.Approver != "ARTEX" || row.Audit.ApproverAddr != "198.51.100.9" {
t.Fatalf("target row = %+v, %v", row, err)
}
// A classic ask goes through the same Decide and records the same.
ctx := WithTaskContext(context.Background(), scope.TaskID, "work#1", nil)
allowed := make(chan bool, 1)
go func() {
allowed <- ic.HandleAsk(ctx, 0, Decision{Action: "ask", Message: "classic ask"}, "Bash", []byte(`{"command":"id"}`))
}()
var classic int64
eventually(t, "the classic ask", func() bool {
return d.QueryRow(`SELECT id FROM intercept_pending WHERE task_id=$1 AND status='pending' AND target_key IS NULL`, scope.TaskID).Scan(&classic) == nil
})
eventually(t, "the classic ask to wait", func() bool { return waiters(ic, classic) == 1 })
if err := ic.Decide(classic, true, "ARTEX", "198.51.100.10"); err != nil {
t.Fatal(err)
}
select {
case ok := <-allowed:
if !ok {
t.Fatal("the approved classic ask was refused")
}
case <-time.After(20 * time.Second):
t.Fatal("the classic ask is still waiting")
}
row, err = d.GetInterceptDetail(classic)
if err != nil || row == nil || row.Status != "allowed" || row.DecidedBy != "human" || row.Audit == nil ||
row.Audit.Approver != "ARTEX" || row.Audit.ApproverAddr != "198.51.100.10" || len(row.Audit.Targets) != 0 {
t.Fatalf("classic row = %+v, %v; want the approver and the address in its audit", row, err)
}
if len(pins(t, d, scope)) != 1 {
t.Fatal("approving a classic ask pinned something")
}
}
// Behind the upstream proxy the proxy chooses the address: the request says the
// approval is name-only and the pin carries the mark. The mark stays when the
// proxy is removed; addresses are enforced again only after a new approval.
func TestPinUnderGlobalProxyIsNameOnly(t *testing.T) {
d := targetTestDB(t)
ic := New(d)
var proxied atomic.Bool
ic.SetUpstreamProxyCheck(proxied.Load)
scope := taskScope(t, d, ic, time.Minute)
const sentence = "This approval is name-only."
approve := func(host string) (*db.TargetPin, string) {
t.Helper()
call := request(context.Background(), ic, scope, nil, target(host, 443, "203.0.113.7"))
id := oneOpenRow(t, d, scope)
row, err := d.GetInterceptDetail(id)
if err != nil || row == nil {
t.Fatalf("request row: %+v, %v", row, err)
}
if err := ic.Decide(id, true, "ARTEX", ""); err != nil {
t.Fatal(err)
}
if res := result(t, call); res.Code != "" {
t.Fatalf("approved call: %+v", res)
}
active, err := d.ActiveTargetPins(scope.dbRef(t), host, 443)
if err != nil || len(active) != 1 {
t.Fatalf("active pins of %s = %+v, %v", host, active, err)
}
return active[0], row.Reason
}
direct, reason := approve("direct.example.com")
if direct.NameOnlyEffective || strings.Contains(reason, sentence) {
t.Fatalf("without a proxy: pin=%+v reason=%q, want an address pin", direct, reason)
}
proxied.Store(true)
named, reason := approve("proxied.example.com")
if !named.NameOnlyEffective || !strings.Contains(reason, sentence) || !strings.Contains(reason, "upstream proxy") {
t.Fatalf("behind the proxy: pin=%+v reason=%q, want a name-only pin and a request that says so", named, reason)
}
if !slices.Equal(named.PinnedIPs, []string{"203.0.113.7"}) {
t.Fatalf("name-only pin addresses = %v, want what the platform resolved", named.PinnedIPs)
}
// The proxy is removed: the pin stays name-only, the address pin stays one.
proxied.Store(false)
ref := scope.dbRef(t)
if still, err := d.ActiveTargetPins(ref, "proxied.example.com", 443); err != nil || len(still) != 1 || !still[0].NameOnlyEffective {
t.Fatalf("name-only pin after the proxy was removed = %+v, %v; want still name-only", still, err)
}
if still, err := d.ActiveTargetPins(ref, "direct.example.com", 443); err != nil || len(still) != 1 || still[0].NameOnlyEffective {
t.Fatalf("address pin = %+v, %v", still, err)
}
// Enforcing addresses again takes a revocation and a new approval.
if _, err := d.RevokeTargetPin(ref, named.ID); err != nil {
t.Fatal(err)
}
if again, _ := approve("proxied.example.com"); again.NameOnlyEffective || again.Version != 2 {
t.Fatalf("pin approved after the proxy was removed = %+v, want an address pin, version 2", again)
}
}
// A new_ip request shows the old version, the pinned addresses and the new
// ones, and its approval extends the pin.
func TestNewIPRequestShowsOldAndNewAddresses(t *testing.T) {
d := targetTestDB(t)
ic := New(d)
scope := taskScope(t, d, ic, time.Minute)
first := request(context.Background(), ic, scope, nil, target("app.example.com", 0, "203.0.113.7"))
id := oneOpenRow(t, d, scope)
if row, err := d.GetInterceptDetail(id); err != nil || row == nil || !strings.Contains(row.Reason, "app.example.com (port not declared: every port)") {
t.Fatalf("port-0 request = %+v, %v; want the reason to say the port is not declared", row, err)
}
if err := ic.Decide(id, true, "ARTEX", ""); err != nil {
t.Fatal(err)
}
if res := result(t, first); res.Code != "" {
t.Fatalf("first call: %+v", res)
}
v1 := pins(t, d, scope)[0]
moved := target("app.example.com", 0, "203.0.113.99")
moved.Reason, moved.Pin = "new_ip", v1
second := request(context.Background(), ic, scope, nil, moved)
id = oneOpenRow(t, d, scope)
row, err := d.GetInterceptDetail(id)
if err != nil || row == nil || !strings.Contains(row.Reason, "pin v1: 203.0.113.7") || !strings.Contains(row.Reason, "now resolves to 203.0.113.99") {
t.Fatalf("new_ip request = %+v, %v; want the old version, the pinned and the new addresses", row, err)
}
if got := row.Audit.Targets[0]; got.Reason != "new_ip" || got.PinID != v1.ID || got.PinVersion != 1 || !slices.Equal(got.PinnedIPs, []string{"203.0.113.7"}) {
t.Fatalf("recorded target = %+v", got)
}
if err := ic.Decide(id, true, "ARTEX", ""); err != nil {
t.Fatal(err)
}
if res := result(t, second); res.Code != "" {
t.Fatalf("second call: %+v", res)
}
active, err := d.ActiveTargetPins(scope.dbRef(t), "app.example.com", 0)
if err != nil || len(active) != 1 || active[0].Version != 2 || !slices.Equal(active[0].PinnedIPs, []string{"203.0.113.7", "203.0.113.99"}) {
t.Fatalf("pin after the new_ip approval = %+v, %v; want v2 with both addresses", active, err)
}
}
// One call, several targets: one request per host and port, and the call
// passes only when all of them are approved.
func TestSeveralTargetsNeedEveryApproval(t *testing.T) {
d := targetTestDB(t)
ic := New(d)
scope := taskScope(t, d, ic, time.Minute)
ask := askRule(t, d, false, 0, "deny")
call := request(context.Background(), ic, scope, &ask,
target("a.example.com", 443, "203.0.113.1"), target("b.example.com", 443, "203.0.113.2"))
var ids []int64
eventually(t, "two open requests", func() bool {
ids = openRows(t, d, scope)
return len(ids) == 2
})
// The ask is merged into the first row the call created; the other is plain.
for n, id := range ids {
row, err := d.GetInterceptDetail(id)
if err != nil || row == nil || row.Audit.MergedAsk != (n == 0) || (row.RuleID != nil) != (n == 0) {
t.Fatalf("row %d = %+v, %v; want the ask merged into the first row only", n, row, err)
}
}
if err := ic.Decide(ids[0], true, "ARTEX", ""); err != nil {
t.Fatal(err)
}
stillWaiting(t, call)
if err := ic.Decide(ids[1], false, "ARTEX", ""); err != nil {
t.Fatal(err)
}
res := result(t, call)
if res.Code != TargetDenied || res.Host != "b.example.com" || !res.AskMerged {
t.Fatalf("call with one target denied: %+v, want %s on b.example.com", res, TargetDenied)
}
if got := pins(t, d, scope); len(got) != 1 || got[0].Host != "a.example.com" {
t.Fatalf("pins = %+v, want only the approved host", got)
}
}
// A scope that names no task and no conversation, or a task id that is not a
// number, never reaches the store.
func TestRequestTargetsRejectsBadScope(t *testing.T) {
ic := New(nil)
for name, scope := range map[string]TargetScope{
"empty": {WaitBound: time.Minute},
"both": {TaskID: "7", ConversationID: 9, WaitBound: time.Minute},
"task not int": {TaskID: "seven", WaitBound: time.Minute},
} {
res := ic.RequestTargets(context.Background(), scope, "WebFetch", webFetchInput, []TargetRequest{target("app.example.com", 443, "203.0.113.7")}, nil)
if res.Code != TargetGateUnavailable {
t.Errorf("%s scope: %+v, want %s", name, res, TargetGateUnavailable)
}
}
}
+17
View File
@@ -0,0 +1,17 @@
package intercept
import (
"os"
"testing"
"time"
"github.com/Autumn-27/artex/db"
)
func TestMain(m *testing.M) {
// Own database per package (#80), as in db, agent, server and evidence.
db.SuffixTestDatabase("intercept")
// A short poll keeps the cross-instance tests fast; production uses 2 s.
targetPollInterval = 50 * time.Millisecond
os.Exit(m.Run())
}
+22
View File
@@ -178,6 +178,28 @@ func verifyJWT(tokenStr string, key []byte) bool {
return err == nil && t.Valid
}
// jwtSubject returns the subject of tokenStr when it is a valid, non-expired
// HS256 token, and "" otherwise.
func jwtSubject(tokenStr string, key []byte) string {
var claims jwt.RegisteredClaims
t, err := jwt.ParseWithClaims(tokenStr, &claims, func(t *jwt.Token) (any, error) {
if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok {
return nil, fmt.Errorf("unexpected signing method")
}
return key, nil
})
if err != nil || !t.Valid {
return ""
}
return claims.Subject
}
// approver is who a human decision is recorded for: the subject of the
// request's session token.
func (s *Server) approver(r *http.Request) string {
return jwtSubject(extractToken(r), s.jwtKey)
}
// extractToken reads the JWT from Authorization: Bearer header,
// artex_token cookie, or ?token= query param (for SSE connections).
func extractToken(r *http.Request) string {
+13 -8
View File
@@ -79,6 +79,7 @@ func (s *Server) pgCreateConversation(w http.ResponseWriter, r *http.Request) {
AgentKey string `json:"agent_key"`
Title string `json:"title"`
LLMProfileID *int64 `json:"llm_profile_id"`
TargetGate bool `json:"target_gate"` // approval gate for asset targets, opt-in
}
if !decodeConversationRequest(w, r, &req) {
return
@@ -115,7 +116,7 @@ func (s *Server) pgCreateConversation(w http.ResponseWriter, r *http.Request) {
writeErr(w, 400, fmt.Sprintf("Title must be at most %d characters", maxConversationTitleRunes))
return
}
c, err := pg.CreateConversation(req.AgentKey, title, req.LLMProfileID)
c, err := pg.CreateConversationGated(req.AgentKey, title, req.LLMProfileID, req.TargetGate)
if err != nil {
writeErr(w, 500, err.Error())
return
@@ -176,14 +177,15 @@ func (s *Server) pgRenameConversation(w http.ResponseWriter, r *http.Request) {
return
}
var req struct {
Title *string `json:"title"`
Pinned *bool `json:"pinned"`
Title *string `json:"title"`
Pinned *bool `json:"pinned"`
TargetGate *bool `json:"target_gate"`
}
if !decodeConversationRequest(w, r, &req) {
return
}
if req.Title == nil && req.Pinned == nil {
writeErr(w, 400, "Provide at least title or pinned")
if req.Title == nil && req.Pinned == nil && req.TargetGate == nil {
writeErr(w, 400, "Provide at least title, pinned or target_gate")
return
}
if req.Title != nil {
@@ -198,7 +200,7 @@ func (s *Server) pgRenameConversation(w http.ResponseWriter, r *http.Request) {
}
req.Title = &title
}
updated, err := pg.UpdateConversation(c.ID, db.ConversationPatch{Title: req.Title, Pinned: req.Pinned})
updated, err := pg.UpdateConversation(c.ID, db.ConversationPatch{Title: req.Title, Pinned: req.Pinned, TargetGate: req.TargetGate})
if err != nil {
writeErr(w, 500, err.Error())
return
@@ -559,8 +561,11 @@ func (s *Server) runConversationTurn(ctx context.Context, cancel context.CancelC
}
}
// Approval scope of this turn's tool calls for the target gate (chat and retester).
// Conversations have no opt-in flag yet, so Enabled stays nil: the gate is off.
gctx := intercept.WithTargetScope(ctx, intercept.TargetScope{ConversationID: c.ID, WaitBound: intercept.DefaultTargetWaitBound})
// The conversation's opt-in flag is the one read when the turn started: a
// switch through the API applies from the next turn.
gate := c.TargetGate
gctx := intercept.WithTargetScope(ctx, intercept.TargetScope{ConversationID: c.ID, Enabled: func() bool { return gate },
WaitBound: intercept.DefaultTargetWaitBound, Interceptor: s.m.interceptor})
// On a manual stop ctx is cancelled; Chat already emits a clean "stopped by the user"
// step, so skip the raw-error entry — only surface genuine failures.
if _, err := ca.Chat(gctx, c.AgentKey, sessionID, msg, maxTurns, maxDuration, webSearch, emit); err != nil {
+2
View File
@@ -63,6 +63,7 @@ type TaskDTO struct {
CompanyIDs []int64 `json:"company_ids"`
CoverageEnabled bool `json:"coverage_enabled"` // asset coverage feature switch (fixed at creation)
PlanQuietMode bool `json:"plan_quiet_mode"` // planner quiet-mode switch (#86), editable mid-run
TargetGate bool `json:"target_gate"` // approval gate for asset targets, opt-in, editable mid-run
}
// FindingSeverityDTO is the per-severity finding count shown in the task list (critical/high/medium/low).
@@ -138,6 +139,7 @@ func taskDTO(t *Task, status string) TaskDTO {
CompanyIDs: lifecycle.CompanyIDs,
CoverageEnabled: t.CoverageEnabled,
PlanQuietMode: t.PlanQuietMode,
TargetGate: t.targetGate.Load(),
}
}
+83 -3
View File
@@ -328,9 +328,9 @@ func interceptFilterParams(q url.Values) (db.InterceptApprovalFilter, error) {
return filter, fmt.Errorf("status must be pending, allowed, denied or timeout")
}
switch filter.DecisionSource {
case "", "model", "rule", "unknown":
case "", "model", "rule", "target", "unknown":
default:
return filter, fmt.Errorf("decision_source must be model, rule or unknown")
return filter, fmt.Errorf("decision_source must be model, rule, target or unknown")
}
switch filter.DecidedBy {
case "", "human", "rule", "model", "system":
@@ -372,7 +372,9 @@ func (s *Server) interceptDecide(w http.ResponseWriter, r *http.Request) {
writeErr(w, 400, "decision must be allowed or denied")
return
}
if err := s.m.interceptor.Decide(id, req.Decision == "allowed"); err != nil {
// On a target request "allowed" pins the addresses the request displayed.
// Every human decision records who took it and from where.
if err := s.m.interceptor.Decide(id, req.Decision == "allowed", s.approver(r), clientIP(r, s.trustedProxies)); err != nil {
if errors.Is(err, intercept.ErrAlreadyDecided) {
writeErr(w, 409, err.Error())
return
@@ -383,6 +385,84 @@ func (s *Server) interceptDecide(w http.ResponseWriter, r *http.Request) {
writeJSON(w, 200, map[string]any{"ok": true})
}
// --- target pins (hosts a human approved for one task or one conversation) ---
// taskPinScope resolves {id} to the pin scope of an existing task.
func (s *Server) taskPinScope(w http.ResponseWriter, r *http.Request) (db.TargetScopeRef, bool) {
id, ok := pathInt(r, "id")
if !ok || id <= 0 {
writeErr(w, 400, "bad task id")
return db.TargetScopeRef{}, false
}
if _, ok := s.m.Task(strconv.FormatInt(id, 10)); !ok {
writeErr(w, 404, "task not found")
return db.TargetScopeRef{}, false
}
return db.TargetScopeRef{TaskID: id}, true
}
func (s *Server) taskListTargetPins(w http.ResponseWriter, r *http.Request) {
if scope, ok := s.taskPinScope(w, r); ok {
s.listTargetPins(w, scope)
}
}
func (s *Server) taskRevokeTargetPin(w http.ResponseWriter, r *http.Request) {
if scope, ok := s.taskPinScope(w, r); ok {
s.revokeTargetPin(w, r, scope)
}
}
func (s *Server) convListTargetPins(w http.ResponseWriter, r *http.Request) {
if _, c, ok := s.convByID(w, r); ok {
s.listTargetPins(w, db.TargetScopeRef{ConversationID: c.ID})
}
}
func (s *Server) convRevokeTargetPin(w http.ResponseWriter, r *http.Request) {
if _, c, ok := s.convByID(w, r); ok {
s.revokeTargetPin(w, r, db.TargetScopeRef{ConversationID: c.ID})
}
}
// listTargetPins returns every pin of the scope, in any state, newest first.
func (s *Server) listTargetPins(w http.ResponseWriter, scope db.TargetScopeRef) {
pg := s.pg(w)
if pg == nil {
return
}
pins, err := pg.ListTargetPins(scope)
if err != nil {
writeErr(w, 500, err.Error())
return
}
writeJSON(w, 200, map[string]any{"pins": pins})
}
// revokeTargetPin revokes pin {pid} of the scope. The next call to that host
// asks again. Revoking a pin that is no longer active changes nothing.
func (s *Server) revokeTargetPin(w http.ResponseWriter, r *http.Request, scope db.TargetScopeRef) {
pg := s.pg(w)
if pg == nil {
return
}
pid, ok := pathInt(r, "pid")
if !ok {
writeErr(w, 400, "bad pin id")
return
}
pin, err := pg.RevokeTargetPin(scope, pid)
if err != nil {
writeErr(w, 500, err.Error())
return
}
if pin == nil {
writeErr(w, 404, "pin not found")
return
}
writeJSON(w, 200, pin)
}
// --- tool-config (global tool interception scope) ---
// interceptGetToolConfig returns the list of tool names that are currently
+16 -4
View File
@@ -13,6 +13,7 @@ import (
"strconv"
"strings"
"sync"
"sync/atomic"
"time"
"github.com/Autumn-27/artex/agent"
@@ -72,6 +73,10 @@ type Task struct {
notify chan struct{}
lifecycleMu sync.RWMutex
llmMu sync.RWMutex
// targetGate is the task's opt-in to the approval gate for asset targets.
// Atomic: it is switched through the API while tool calls read it.
targetGate atomic.Bool
interceptor *intercept.Interceptor // owner of the target gate, handed to the scope
// pendingTriggers accumulates the concrete changes (worker done / finding) that
// fired planning rounds since the last one consumed them. The debounce coalesces
@@ -397,6 +402,8 @@ func NewManager(dir, proxyAddr string) (*Manager, error) {
log.Printf("[llmusage] create table: %v", err)
}
m := &Manager{dir: dir, pg: pg, assets: pg.Assets(), tasks: map[string]*Task{}, interceptor: intercept.New(pg)}
// Behind the global egress proxy a target approval is name-only.
m.interceptor.SetUpstreamProxyCheck(func() bool { return m.GlobalProxy() != "" })
if proxyAddr != "" {
tr, err := traffic.Open(filepath.Join(dir, "traffic"), proxyAddr)
if err != nil {
@@ -852,7 +859,7 @@ func unixNanoOrZero(t *time.Time) int64 {
}
func taskFromPG(pt *pgdb.Task, store *pgdb.ExplorationStore, ic *intercept.Interceptor) *Task {
return &Task{
t := &Task{
ID: strconv.FormatInt(pt.ID, 10), ExpID: pt.ExplorationID,
Name: pt.Name,
CategoryID: cloneInt64Ptr(pt.CategoryID), CategoryName: pt.CategoryName,
@@ -871,14 +878,18 @@ func taskFromPG(pt *pgdb.Task, store *pgdb.ExplorationStore, ic *intercept.Inter
CoverageEnabled: pt.CoverageEnabled,
FirstRunAt: unixOrZero(pt.FirstRunAt), DeadlineAt: unixOrZero(pt.DeadlineAt),
Store: store, Guard: guard.NewWithInterceptor(ic), TargetGuard: guard.NewTargetOnly(ic), notify: make(chan struct{}, 1),
interceptor: ic,
}
t.targetGate.Store(pt.TargetGate)
return t
}
// targetScope is the approval scope of this task's tool calls for the target
// gate. wait bounds how long a gated call may wait for a human decision. Tasks
// have no opt-in flag yet, so Enabled stays nil: the gate is off.
// gate. wait bounds how long a gated call may wait for a human decision.
// Enabled reads the task's opt-in flag on every call, so a switch through the
// API takes effect on the next tool call.
func (t *Task) targetScope(wait time.Duration) intercept.TargetScope {
return intercept.TargetScope{TaskID: t.ID, WaitBound: wait}
return intercept.TargetScope{TaskID: t.ID, Enabled: t.targetGate.Load, WaitBound: wait, Interceptor: t.interceptor}
}
// UpdateTaskMetadata changes list-only task metadata without interrupting any
@@ -907,6 +918,7 @@ func (m *Manager) UpdateTaskMetadata(taskID string, patch pgdb.TaskPatch) (*Task
// Live switch (#86): runRound reads the flag per round, so the next
// planning round (edge or heartbeat) already respects the new value.
task.PlanQuietMode = updated.PlanQuietMode
task.targetGate.Store(updated.TargetGate)
return task, nil
}
+6
View File
@@ -715,6 +715,8 @@ func (s *Server) Handler() http.Handler {
mux.HandleFunc("PUT /api/tasks/{id}/intercept-rules/{rid}", s.taskInterceptUpdateRule)
mux.HandleFunc("DELETE /api/tasks/{id}/intercept-rules/{rid}", s.taskInterceptDeleteRule)
mux.HandleFunc("POST /api/tasks/{id}/intercept-rules/{rid}/toggle", s.taskInterceptToggleRule)
mux.HandleFunc("GET /api/tasks/{id}/target-pins", s.taskListTargetPins)
mux.HandleFunc("POST /api/tasks/{id}/target-pins/{pid}/revoke", s.taskRevokeTargetPin)
mux.HandleFunc("POST /api/tasks/control/batch", s.controlTasksBatch)
mux.HandleFunc("GET /api/task-archives", s.listTaskArchives)
mux.HandleFunc("GET /api/task-archives/{id}", s.getTaskArchive)
@@ -860,6 +862,8 @@ func (s *Server) Handler() http.Handler {
mux.HandleFunc("POST /api/conversations/{id}/messages", s.pgSendConversationMessage)
mux.HandleFunc("POST /api/conversations/{id}/stop", s.pgStopConversation)
mux.HandleFunc("GET /api/conversations/{id}/messages/{seq}", s.pgConversationMsgDetail)
mux.HandleFunc("GET /api/conversations/{id}/target-pins", s.convListTargetPins)
mux.HandleFunc("POST /api/conversations/{id}/target-pins/{pid}/revoke", s.convRevokeTargetPin)
mux.HandleFunc("GET /api/agents", s.pgListAgents)
mux.HandleFunc("POST /api/agents", s.pgCreateAgent)
@@ -1584,6 +1588,7 @@ type createTaskReq struct {
LLMRotate *bool `json:"llm_rotate"` // task-level timeout (seconds); 0/omitted = no limit
PlanHeartbeatSeconds int `json:"plan_heartbeat_seconds"` // planner heartbeat interval (seconds); 0/omitted = default 600 (10min); floor = default = 600, lower values are raised to 600
PlanQuietMode *bool `json:"plan_quiet_mode,omitempty"` // planner quiet-mode switch (#86); omitted = default on; false = edge rounds always plan
TargetGate *bool `json:"target_gate,omitempty"` // approval gate for asset targets, opt-in; omitted/null = off
SeedFirstIntent *bool `json:"seed_first_intent,omitempty"` // issue one seed intent at creation (content = description + goal) so a worker can start right away without waiting for the first planner round; omitted/null = default off, the standard plan-then-execute flow. Pass true explicitly to enable (for CTF tasks often solved in one work, this skips the pre-start planner round).
CoverageEnabled *bool `json:"coverage_enabled,omitempty"` // asset coverage feature; omitted/null = default on (true). false = disables coverage calculation/display/auto-accumulated scope + hides add_task_scope/list_untested_assets. company links are unaffected.
// InterceptRules: task-level asset interception/allow rules (entered at creation, stored in task_intercept_rules, never the global table).
@@ -1660,6 +1665,7 @@ func (s *Server) createTask(w http.ResponseWriter, r *http.Request) {
CoverageEnabled: req.CoverageEnabled,
InterceptRules: interceptRules,
LLMRotate: req.LLMRotate,
TargetGate: req.TargetGate,
})
if err != nil {
if errors.Is(err, db.ErrTaskCategoryInvalid) || errors.Is(err, db.ErrTaskCategoryNotFound) {
+553
View File
@@ -0,0 +1,553 @@
package server
import (
"context"
"encoding/json"
"errors"
"fmt"
"net/http/httptest"
"slices"
"strconv"
"strings"
"testing"
"time"
"github.com/Autumn-27/artex/agent"
"github.com/Autumn-27/artex/db"
"github.com/Autumn-27/artex/intercept"
)
// Target-gate opt-in flag, pins and decisions over the real mux and the real
// auth. The server context is already cancelled, so a task created through
// these tests never starts an engine.
type gateServer struct {
s *Server
m *Manager
token string
}
func newGateServer(t *testing.T) *gateServer {
t.Helper()
skipWithoutDSN(t)
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = m.Close() })
ctx, cancel := context.WithCancel(context.Background())
cancel()
s := New(ctx, m, t.TempDir(), t.TempDir(), t.TempDir())
token, err := signJWT(s.jwtKey)
if err != nil {
t.Fatal(err)
}
return &gateServer{s: s, m: m, token: token}
}
// do sends an authenticated request. httptest sets the peer to 192.0.2.1.
func (g *gateServer) do(t *testing.T, method, path, body string) (int, map[string]any) {
t.Helper()
req := httptest.NewRequest(method, path, strings.NewReader(body))
req.Header.Set("Authorization", "Bearer "+g.token)
rec := httptest.NewRecorder()
g.s.Handler().ServeHTTP(rec, req)
var out map[string]any
_ = json.Unmarshal(rec.Body.Bytes(), &out)
return rec.Code, out
}
// newTask creates a task through the manager and removes it, its pins and its
// approval rows when the test ends.
func (g *gateServer) newTask(t *testing.T, opts db.TaskCreateOptions) *Task {
t.Helper()
created, err := g.m.CreateTaskWithOptions("target gate flag", "flag, pins and decisions", opts)
if err != nil {
t.Fatal(err)
}
g.cleanupTask(t, created.ID)
task, ok := g.m.Task(created.ID)
if !ok {
t.Fatal("task not found after creation")
}
return task
}
func (g *gateServer) cleanupTask(t *testing.T, id string) {
t.Helper()
t.Cleanup(func() {
n, _ := strconv.ParseInt(id, 10, 64)
_ = g.m.pg.DeleteTask(n)
_, _ = g.m.pg.Exec(`DELETE FROM target_pins WHERE task_id=$1`, n)
_, _ = g.m.pg.Exec(`DELETE FROM intercept_pending WHERE task_id=$1`, id)
})
}
func taskRef(t *testing.T, task *Task) db.TargetScopeRef {
t.Helper()
id, err := strconv.ParseInt(task.ID, 10, 64)
if err != nil {
t.Fatal(err)
}
return db.TargetScopeRef{TaskID: id}
}
// openTargetRow opens a target approval request showing ips.
func openTargetRow(t *testing.T, pg *db.DB, scope db.TargetScopeRef, host string, ips ...string) int64 {
t.Helper()
id, created, err := pg.FindOrCreateTargetRequest(db.TargetRequest{
Scope: scope, Host: host, Port: 443, AgentName: "worker", ToolName: "WebFetch", Reason: "target " + host + " is not approved",
Audit: &db.InterceptAudit{Correlation: "unavailable", InitialAction: "ask", ExecutionStatus: "not_started",
Targets: []db.InterceptTarget{{Host: host, Port: 443, ResolvedIPs: ips, Resolver: "system", ResolvedAt: time.Now().UTC(), Outcome: "requested"}}},
})
if err != nil || !created {
t.Fatalf("open target request: id=%d created=%v err=%v", id, created, err)
}
return id
}
// The flag is read on every tool call: a switch through the API is seen by a
// scope that was handed out before it.
func TestTargetGateLiveSwitch(t *testing.T) {
g := newGateServer(t)
task := g.newTask(t, db.TaskCreateOptions{})
scope := task.targetScope(intercept.DefaultTargetWaitBound)
if scope.Enabled == nil || scope.Enabled() {
t.Fatal("a new task must have the gate off")
}
if scope.Interceptor != g.m.interceptor {
t.Fatal("the task scope does not carry the manager's interceptor")
}
if code, body := g.do(t, "GET", "/api/tasks/"+task.ID, ""); code != 200 || body["target_gate"] != false {
t.Fatalf("task detail: %d %v, want target_gate false", code, body)
}
code, body := g.do(t, "PATCH", "/api/tasks/"+task.ID, `{"target_gate":true}`)
if code != 200 || body["target_gate"] != true {
t.Fatalf("switch on: %d %v", code, body)
}
if !scope.Enabled() {
t.Fatal("the scope handed out before the switch still reports the gate off")
}
id, _ := strconv.ParseInt(task.ID, 10, 64)
if stored, err := g.m.pg.GetTask(id); err != nil || stored == nil || !stored.TargetGate {
t.Fatalf("stored task = %+v, %v; want the gate on", stored, err)
}
// Another metadata patch leaves the flag alone.
if code, body := g.do(t, "PATCH", "/api/tasks/"+task.ID, `{"name":"renamed"}`); code != 200 || body["target_gate"] != true || !scope.Enabled() {
t.Fatalf("rename: %d %v, want the gate still on", code, body)
}
if code, body := g.do(t, "PATCH", "/api/tasks/"+task.ID, `{"target_gate":false}`); code != 200 || body["target_gate"] != false || scope.Enabled() {
t.Fatalf("switch off: %d %v enabled=%v", code, body, scope.Enabled())
}
if code, body := g.do(t, "PATCH", "/api/tasks/"+task.ID, `{}`); code != 400 || !strings.Contains(fmt.Sprint(body["error"]), "target_gate") {
t.Fatalf("empty patch: %d %v, want 400 naming target_gate", code, body)
}
}
const gateTestGoal = "Probe the lab target at http://127.0.0.1:9/v1 for unauthenticated data exposure; record one record as evidence and stop."
// The flag is set on the existing create path. Omitted means off, and a task
// loaded from the database keeps what was stored.
func TestCreateTaskAcceptsTargetGate(t *testing.T) {
g := newGateServer(t)
for _, tc := range []struct {
field string
want bool
}{
{`,"target_gate":true`, true},
{`,"target_gate":false`, false},
{``, false},
} {
code, body := g.do(t, "POST", "/api/tasks", `{"description":"gate create","goal":"`+gateTestGoal+`"`+tc.field+`}`)
id, _ := body["id"].(string)
if code != 201 || id == "" {
t.Fatalf("create %q: %d %v", tc.field, code, body)
}
g.cleanupTask(t, id)
task, ok := g.m.Task(id)
if !ok || body["target_gate"] != tc.want || task.targetScope(0).Enabled() != tc.want {
t.Fatalf("create %q: target_gate=%v, want %v", tc.field, body["target_gate"], tc.want)
}
n, _ := strconv.ParseInt(id, 10, 64)
stored, err := g.m.pg.GetTask(n)
if err != nil || stored == nil || stored.TargetGate != tc.want {
t.Fatalf("create %q: stored %+v, %v; want %v", tc.field, stored, err, tc.want)
}
if loaded := taskFromPG(stored, g.m.pg.Exploration(stored.ExplorationID), g.m.interceptor); loaded.targetScope(0).Enabled() != tc.want {
t.Fatalf("create %q: a task loaded from the database has the gate %v", tc.field, !tc.want)
}
}
}
// No inheritance: a task spawned by the orchestration agent starts with the
// gate off, even from a gated parent.
func TestSpawnedTaskStartsWithGateOff(t *testing.T) {
g := newGateServer(t)
on := true
parent := g.newTask(t, db.TaskCreateOptions{TargetGate: &on})
if !parent.targetScope(0).Enabled() {
t.Fatal("control: the parent was expected to have the gate on")
}
in, _ := json.Marshal(map[string]any{"description": "spawned child", "goal": gateTestGoal, "parent_ref": parent.ID, "source_task_ids": []string{parent.ID}})
res, err := g.s.toolSpawnTask().Call(context.Background(), in, nil)
if err != nil {
t.Fatal(err)
}
childID, found := strings.CutPrefix(res.Flatten(), "task created: ")
if !found {
t.Fatalf("spawn_task result: %q", res.Flatten())
}
g.cleanupTask(t, childID)
child, ok := g.m.Task(childID)
if !ok || child.ParentRef != parent.ID {
t.Fatalf("spawned task = %+v (found=%v), want a child of %s", child, ok, parent.ID)
}
if child.targetScope(0).Enabled() {
t.Fatal("the spawned task inherited the gate")
}
n, _ := strconv.ParseInt(childID, 10, 64)
if stored, err := g.m.pg.GetTask(n); err != nil || stored == nil || stored.TargetGate {
t.Fatalf("stored child = %+v, %v; want the gate off", stored, err)
}
if pins, err := g.m.pg.ListTargetPins(db.TargetScopeRef{TaskID: n}); err != nil || len(pins) != 0 {
t.Fatalf("the spawned task has pins: %+v, %v", pins, err)
}
}
// The conversation flag is set on create and on patch, and the turn that starts
// after a switch runs with the new value.
func TestConversationTargetGateOverHTTP(t *testing.T) {
g := newGateServer(t)
turnScope := func(id int64) intercept.TargetScope {
t.Helper()
conv, err := g.m.pg.GetConversation(id)
if err != nil || conv == nil {
t.Fatalf("conversation %d: %+v, %v", id, conv, err)
}
probe := newScopeProbeProvider()
g.s.cfgMu.Lock()
g.s.chatAgent = agent.NewChatAgent(probe, "test-model", g.m.dir, nil, 10000)
g.s.cfgMu.Unlock()
// The server context of this fixture is cancelled; the turn gets its own.
ctx, cancel := context.WithCancelCause(intercept.WithConvID(context.Background(), id))
g.s.runConversationTurn(ctx, cancel, conv, "hello", g.s.convBusyKey(id))
got := probe.first(t)
if !got.ok || got.scope.ConversationID != id || got.scope.Enabled == nil {
t.Fatalf("turn scope = %+v (present=%v)", got.scope, got.ok)
}
return got.scope
}
create := func(body string) int64 {
t.Helper()
code, out := g.do(t, "POST", "/api/conversations", body)
id, _ := out["id"].(float64)
if code != 200 || id == 0 {
t.Fatalf("create conversation: %d %v", code, out)
}
t.Cleanup(func() { _ = g.m.pg.DeleteConversation(int64(id)) })
return int64(id)
}
plain := create(`{"agent_key":"auto","title":"gate off"}`)
if scope := turnScope(plain); scope.Enabled() {
t.Fatal("a conversation created without the flag runs with the gate on")
}
gated := create(`{"agent_key":"auto","title":"gate on","target_gate":true}`)
scope := turnScope(gated)
if !scope.Enabled() || scope.Interceptor != g.m.interceptor {
t.Fatalf("gated conversation: enabled=%v, want the gate on and the manager's interceptor", scope.Enabled())
}
path := "/api/conversations/" + strconv.FormatInt(gated, 10)
if code, out := g.do(t, "PATCH", path, `{"target_gate":false}`); code != 200 || out["target_gate"] != false {
t.Fatalf("switch off: %d %v", code, out)
}
if turnScope(gated).Enabled() {
t.Fatal("the turn after the switch still runs with the gate on")
}
if code, out := g.do(t, "PATCH", path, `{"title":"renamed","target_gate":true}`); code != 200 || out["target_gate"] != true || out["title"] != "renamed" {
t.Fatalf("switch on with a rename: %d %v", code, out)
}
if code, out := g.do(t, "PATCH", path, `{}`); code != 400 || !strings.Contains(fmt.Sprint(out["error"]), "target_gate") {
t.Fatalf("empty patch: %d %v, want 400 naming target_gate", code, out)
}
}
// The decide endpoint keeps its request and response shape. On a target row
// "allowed" pins the displayed addresses. Every human decision, on a target row
// or on a classic ask, records the subject of the session token and the remote
// address; decided_by stays the actor class.
func TestInterceptDecideRecordsApproverOverHTTP(t *testing.T) {
g := newGateServer(t)
task := g.newTask(t, db.TaskCreateOptions{})
scope := taskRef(t, task)
pg := g.m.pg
decide := func(id int64, decision string) int {
t.Helper()
code, _ := g.do(t, "POST", fmt.Sprintf("/api/intercept/pending/%d/decide", id), `{"decision":"`+decision+`"}`)
return code
}
target := openTargetRow(t, pg, scope, "app.example.com", "203.0.113.7")
if code := decide(target, "allowed"); code != 200 {
t.Fatalf("decide target row: %d", code)
}
pins, err := pg.ListTargetPins(scope)
if err != nil || len(pins) != 1 || pins[0].DecidedBy != "ARTEX" || pins[0].DecidedFrom != "192.0.2.1" ||
!slices.Equal(pins[0].PinnedIPs, []string{"203.0.113.7"}) || pins[0].InterceptID == nil || *pins[0].InterceptID != target {
t.Fatalf("pins = %+v, %v; want one pin with the approver and the address", pins, err)
}
row, err := pg.GetInterceptDetail(target)
if err != nil || row == nil || row.Status != "allowed" || row.DecidedBy != "human" || row.Audit.Approver != "ARTEX" || row.Audit.ApproverAddr != "192.0.2.1" {
t.Fatalf("target row = %+v, %v", row, err)
}
if code := decide(target, "denied"); code != 409 {
t.Fatalf("second decision: %d, want 409", code)
}
denied := openTargetRow(t, pg, scope, "other.example.com", "203.0.113.8")
if code := decide(denied, "denied"); code != 200 {
t.Fatalf("deny target row: %d", code)
}
if pins, err := pg.ListTargetPins(scope); err != nil || len(pins) != 1 {
t.Fatalf("pins after a denial = %+v, %v; want still one", pins, err)
}
classic, err := pg.CreateInterceptPending(0, 0, task.ID, "worker", "Bash", []byte(`{"command":"id"}`), "classic ask",
&db.InterceptAudit{Correlation: "exact", InitialAction: "ask", ExecutionStatus: "not_started"})
if err != nil {
t.Fatal(err)
}
if code := decide(classic, "allowed"); code != 200 {
t.Fatalf("decide classic ask: %d", code)
}
row, err = pg.GetInterceptDetail(classic)
if err != nil || row == nil || row.Status != "allowed" || row.DecidedBy != "human" || row.Audit == nil ||
row.Audit.Approver != "ARTEX" || row.Audit.ApproverAddr != "192.0.2.1" || row.Audit.EffectiveAction != "allow" {
t.Fatalf("classic row = %+v, %v; want the approver and the address in its audit", row, err)
}
if pins, err := pg.ListTargetPins(scope); err != nil || len(pins) != 1 {
t.Fatalf("pins after a classic approval = %+v, %v; want still one", pins, err)
}
// Target rows are listed under their own decision source.
code, list := g.do(t, "GET", "/api/intercept/task/"+task.ID+"?decision_source=target&page=1&size=50", "")
if items, _ := list["items"].([]any); code != 200 || len(items) != 2 {
t.Fatalf("decision_source=target: %d %v, want the two target rows", code, list)
}
if code, _ := g.do(t, "GET", "/api/intercept/history?decision_source=nonsense&page=1", ""); code != 400 {
t.Fatalf("unknown decision_source: %d, want 400", code)
}
}
func TestJWTSubject(t *testing.T) {
key := []byte("0123456789abcdef0123456789abcdef")
token, err := signJWT(key)
if err != nil {
t.Fatal(err)
}
if got := jwtSubject(token, key); got != "ARTEX" {
t.Fatalf("subject = %q, want ARTEX", got)
}
if got := jwtSubject(token, []byte("another key, another key, another")); got != "" {
t.Fatalf("subject of a token signed with another key = %q, want none", got)
}
if got := jwtSubject("", key); got != "" {
t.Fatalf("subject of no token = %q", got)
}
}
// Pin routes: list and revoke, scoped to one existing task or conversation.
func TestTargetPinRoutes(t *testing.T) {
g := newGateServer(t)
pg := g.m.pg
task, other := g.newTask(t, db.TaskCreateOptions{}), g.newTask(t, db.TaskCreateOptions{})
conv, err := pg.CreateConversation("auto", "pin routes", nil)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
_ = pg.DeleteConversation(conv.ID)
_, _ = pg.Exec(`DELETE FROM target_pins WHERE conversation_id=$1`, conv.ID)
})
pin := func(scope db.TargetScopeRef, host string) int64 {
t.Helper()
p, ok, err := pg.ResolveTargetRequest(openTargetRow(t, pg, scope, host, "203.0.113.7"), true, "ARTEX", "192.0.2.1")
if err != nil || !ok || p == nil {
t.Fatalf("pin %s: %+v ok=%v err=%v", host, p, ok, err)
}
return p.ID
}
taskPin := pin(taskRef(t, task), "app.example.com")
otherPin := pin(taskRef(t, other), "app.example.com")
convPin := pin(db.TargetScopeRef{ConversationID: conv.ID}, "chat.example.com")
listed := func(path string) []any {
t.Helper()
code, body := g.do(t, "GET", path, "")
pins, ok := body["pins"].([]any)
if code != 200 || !ok {
t.Fatalf("GET %s: %d %v", path, code, body)
}
return pins
}
for _, tc := range []struct {
base string
id int64
host string
}{
{"/api/tasks/" + task.ID, taskPin, "app.example.com"},
{"/api/conversations/" + strconv.FormatInt(conv.ID, 10), convPin, "chat.example.com"},
} {
pins := listed(tc.base + "/target-pins")
if len(pins) != 1 {
t.Fatalf("%s: %d pins, want 1", tc.base, len(pins))
}
got := pins[0].(map[string]any)
if got["id"] != float64(tc.id) || got["host"] != tc.host || got["status"] != "active" || got["decided_by"] != "ARTEX" || got["name_only_effective"] != false {
t.Fatalf("%s: listed pin = %v", tc.base, got)
}
revoke := fmt.Sprintf("%s/target-pins/%d/revoke", tc.base, tc.id)
code, body := g.do(t, "POST", revoke, "")
if code != 200 || body["status"] != "revoked" || body["revoked_reason"] != "operator" {
t.Fatalf("POST %s: %d %v", revoke, code, body)
}
at := body["revoked_at"]
// Revoking again changes nothing.
if code, body := g.do(t, "POST", revoke, ""); code != 200 || body["status"] != "revoked" || body["revoked_at"] != at {
t.Fatalf("second POST %s: %d %v", revoke, code, body)
}
if pins := listed(tc.base + "/target-pins"); len(pins) != 1 || pins[0].(map[string]any)["status"] != "revoked" {
t.Fatalf("%s after the revoke: %v, want the pin kept as revoked", tc.base, pins)
}
// A pin of another scope is not reachable through this one.
if code, _ := g.do(t, "POST", fmt.Sprintf("%s/target-pins/%d/revoke", tc.base, otherPin), ""); code != 404 {
t.Fatalf("revoke of a foreign pin through %s: %d, want 404", tc.base, code)
}
}
if active, err := pg.ActiveTargetPins(taskRef(t, other), "app.example.com", 443); err != nil || len(active) != 1 {
t.Fatalf("the pin of the other task = %+v, %v; want still active", active, err)
}
for _, path := range []string{"/api/tasks/999999999/target-pins", "/api/conversations/999999999/target-pins"} {
if code, _ := g.do(t, "GET", path, ""); code != 404 {
t.Fatalf("GET %s: %d, want 404", path, code)
}
}
if code, _ := g.do(t, "POST", "/api/tasks/"+task.ID+"/target-pins/notanumber/revoke", ""); code != 400 {
t.Fatalf("revoke with a bad pin id: %d, want 400", code)
}
}
// A worker that loses its claim while it waits in the gate stops waiting, and
// the request stays open: it is host-level, not owned by the call. The worker
// that took the intent over attaches to the same request.
func TestClaimLostEndsGateWait(t *testing.T) {
g := newGateServer(t)
task := g.newTask(t, db.TaskCreateOptions{})
pg, ic := g.m.pg, g.m.interceptor
store := task.Store
intent, err := store.AddIntent(map[string]any{"summary": "leased"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
token, err := store.ClaimIntentToken(intent, "work#1")
if err != nil || token == "" {
t.Fatalf("claim: token=%q err=%v", token, err)
}
ctx, cancel := context.WithCancelCause(intercept.WithTaskContext(context.Background(), task.ID, "work#1", nil))
defer cancel(nil)
stop := startClaimHeartbeat(store, 20*time.Millisecond, "work#1", intent, token, cancel)
defer stop()
scope := task.targetScope(time.Minute)
targets := []intercept.TargetRequest{{Host: "app.example.com", Port: 443, ResolvedIPs: []string{"203.0.113.7"}, Resolver: "system", ResolvedAt: time.Now().UTC(), Reason: "no_pin"}}
input := []byte(`{"url":"https://app.example.com/"}`)
stale := make(chan intercept.TargetResult, 1)
go func() { stale <- ic.RequestTargets(ctx, scope, "WebFetch", input, targets, nil) }()
openRows := func() []int64 {
t.Helper()
rows, err := pg.Query(`SELECT id FROM intercept_pending WHERE task_id=$1 AND status='pending' AND target_key IS NOT NULL ORDER BY id`, task.ID)
if err != nil {
t.Fatal(err)
}
defer rows.Close()
var ids []int64
for rows.Next() {
var id int64
if err := rows.Scan(&id); err != nil {
t.Fatal(err)
}
ids = append(ids, id)
}
return ids
}
waitUntil(t, "the gate request", func() bool { return len(openRows()) == 1 })
request := openRows()[0]
select {
case res := <-stale:
t.Fatalf("the call returned %+v while its claim was live", res)
case <-time.After(200 * time.Millisecond):
}
// Takeover, as in TestClaimHeartbeatRenewsThenCancelsOnTakeover.
var reclaimed int64
for attempt := 0; attempt < 200 && reclaimed == 0; attempt++ {
if _, err := pg.Exec(`UPDATE exploration_nodes SET last_heartbeat_at = now() - $2 * INTERVAL '1 second' WHERE id=$1`,
intent, (db.IntentClaimTTL + time.Second).Seconds()); err != nil {
t.Fatal(err)
}
if reclaimed, err = store.ResetRunningIntents(); err != nil {
t.Fatal(err)
}
}
if reclaimed != 1 {
t.Fatalf("could not take the claim over: reclaimed=%d", reclaimed)
}
if other, err := store.ClaimIntentToken(intent, "work#2"); err != nil || other == "" {
t.Fatalf("re-claim: token=%q err=%v", other, err)
}
var res intercept.TargetResult
select {
case res = <-stale:
case <-time.After(10 * time.Second):
t.Fatal("the gate wait did not end after the claim was lost")
}
if cause := context.Cause(ctx); !errors.Is(cause, agent.AbortWorkClaimLost) {
t.Fatalf("cancel cause=%v, want AbortWorkClaimLost", cause)
}
if res.Code != intercept.TargetApprovalPending || res.PendingID != request {
t.Fatalf("stale call: %+v, want %s on request %d", res, intercept.TargetApprovalPending, request)
}
if row, err := pg.GetInterceptPending(request); err != nil || row == nil || row.Status != "pending" || row.DecidedBy != "" {
t.Fatalf("request after the claim was lost = %+v, %v; want still open", row, err)
}
// The second claimer re-runs the intent, hits the same host and attaches.
next := make(chan intercept.TargetResult, 1)
go func() {
next <- ic.RequestTargets(intercept.WithTaskContext(context.Background(), task.ID, "work#2", nil), scope, "WebFetch", input, targets, nil)
}()
select {
case res := <-next:
t.Fatalf("the second claimer returned %+v before a decision", res)
case <-time.After(200 * time.Millisecond):
}
if ids := openRows(); !slices.Equal(ids, []int64{request}) {
t.Fatalf("open requests = %v, want the second claimer on the same request %d", ids, request)
}
if err := ic.Decide(request, true, "ARTEX", "192.0.2.1"); err != nil {
t.Fatal(err)
}
select {
case res := <-next:
if res.Code != "" {
t.Fatalf("second claimer after the approval: %+v", res)
}
case <-time.After(10 * time.Second):
t.Fatal("the second claimer is still waiting after the approval")
}
if pins, err := pg.ListTargetPins(taskRef(t, task)); err != nil || len(pins) != 1 {
t.Fatalf("pins = %+v, %v; want one", pins, err)
}
}
+10 -4
View File
@@ -335,8 +335,11 @@ func TestEveryAgentKindCarriesTargetScope(t *testing.T) {
if got.scope.WaitBound != site.wait {
t.Fatalf("wait bound = %v, want %v", got.scope.WaitBound, site.wait)
}
if got.scope.Enabled != nil {
t.Fatal("the gate must be off: tasks have no opt-in flag yet")
if got.scope.Enabled == nil || got.scope.Enabled() {
t.Fatal("the gate must be off: the task did not opt in")
}
if got.scope.Interceptor != f.m.interceptor {
t.Fatal("the scope does not carry the interceptor that owns the gate")
}
})
}
@@ -348,8 +351,11 @@ func TestEveryAgentKindCarriesTargetScope(t *testing.T) {
if got.scope.WaitBound != intercept.DefaultTargetWaitBound {
t.Fatalf("wait bound = %v, want %v", got.scope.WaitBound, intercept.DefaultTargetWaitBound)
}
if got.scope.Enabled != nil {
t.Fatal("the gate must be off: conversations have no opt-in flag yet")
if got.scope.Enabled == nil || got.scope.Enabled() {
t.Fatal("the gate must be off: the conversation did not opt in")
}
if got.scope.Interceptor != f.m.interceptor {
t.Fatal("the scope does not carry the interceptor that owns the gate")
}
})
}
+4 -3
View File
@@ -28,6 +28,7 @@ func (s *Server) updateTaskMetadata(w http.ResponseWriter, r *http.Request) {
Name *string `json:"name"`
Pinned *bool `json:"pinned"`
PlanQuietMode *bool `json:"plan_quiet_mode"`
TargetGate *bool `json:"target_gate"`
}
if err := decode(r, &request); err != nil {
var tooLarge *http.MaxBytesError
@@ -38,8 +39,8 @@ func (s *Server) updateTaskMetadata(w http.ResponseWriter, r *http.Request) {
}
return
}
if request.Name == nil && request.Pinned == nil && request.PlanQuietMode == nil {
writeErr(w, http.StatusBadRequest, "at least one of name, pinned or plan_quiet_mode must be provided")
if request.Name == nil && request.Pinned == nil && request.PlanQuietMode == nil && request.TargetGate == nil {
writeErr(w, http.StatusBadRequest, "at least one of name, pinned, plan_quiet_mode or target_gate must be provided")
return
}
if request.Name != nil {
@@ -54,7 +55,7 @@ func (s *Server) updateTaskMetadata(w http.ResponseWriter, r *http.Request) {
}
request.Name = &name
}
task, err := s.m.UpdateTaskMetadata(taskID, db.TaskPatch{Name: request.Name, Pinned: request.Pinned, PlanQuietMode: request.PlanQuietMode})
task, err := s.m.UpdateTaskMetadata(taskID, db.TaskPatch{Name: request.Name, Pinned: request.Pinned, PlanQuietMode: request.PlanQuietMode, TargetGate: request.TargetGate})
if err != nil {
writeErr(w, http.StatusInternalServerError, err.Error())
return