mirror of
https://github.com/BaaRai/ARTEX.git
synced 2026-11-15 04:30:55 +08:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
f07e84ac7a | ||
|
|
a77851db7d | ||
|
|
c0b741bcf6 | ||
|
|
ec3db7fdec |
+4
-4
@@ -162,7 +162,7 @@ func (g *Guard) applyIntercept(ctx context.Context, ev hook.Event) hook.Result {
|
||||
if gated {
|
||||
var res intercept.TargetResult
|
||||
if unpinned, res = g.gate.CheckTargets(ctx, scope, ev.ToolName, ev.Input); res.Code != "" {
|
||||
return g.block(ev.ToolName, targetBlockMessage(res), "")
|
||||
return g.block(ev.ToolName, TargetBlockMessage(res), "")
|
||||
}
|
||||
}
|
||||
var ask *intercept.Decision
|
||||
@@ -182,7 +182,7 @@ func (g *Guard) applyIntercept(ctx context.Context, ev hook.Event) hook.Result {
|
||||
if len(unpinned) > 0 {
|
||||
res, merged := g.requestTargets(ctx, scope, ev, unpinned, ask)
|
||||
if res.Code != "" {
|
||||
return g.block(ev.ToolName, targetBlockMessage(res), "")
|
||||
return g.block(ev.ToolName, TargetBlockMessage(res), "")
|
||||
}
|
||||
if merged {
|
||||
// The approval of the target request answered the ask as well.
|
||||
@@ -235,11 +235,11 @@ var targetGuidance = map[string]string{
|
||||
"Retry the same call later, or report the failure.",
|
||||
}
|
||||
|
||||
// targetBlockMessage is systemBlockMessage for the target gate: the block code,
|
||||
// TargetBlockMessage is systemBlockMessage for the target gate: the block code,
|
||||
// the reason and the target, then what to do about it. A target block is not
|
||||
// always final (an approval can still arrive, a lookup can succeed later), so
|
||||
// it does not end with "This operation is forbidden."
|
||||
func targetBlockMessage(res intercept.TargetResult) string {
|
||||
func TargetBlockMessage(res intercept.TargetResult) string {
|
||||
reason := res.Reason
|
||||
if res.Host != "" && !strings.Contains(reason, res.Host) {
|
||||
target := res.Host
|
||||
|
||||
@@ -257,7 +257,7 @@ func TestTargetBlockMessage(t *testing.T) {
|
||||
[]string{"TARGET_GATE_UNAVAILABLE: target app.example.com:443: the pins could not be read: connection refused.", "failure on the platform side"}},
|
||||
}
|
||||
for _, c := range cases {
|
||||
msg := targetBlockMessage(c.res)
|
||||
msg := TargetBlockMessage(c.res)
|
||||
if !strings.HasPrefix(msg, "[ARTEX platform control, not target defense] "+c.res.Code+": ") {
|
||||
t.Errorf("%s: message = %q", c.res.Code, msg)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,154 @@
|
||||
package intercept
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/url"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// The custom tools render their target from a template and the parameters of
|
||||
// the call, after the hook has seen the call: the hook reads the parameters, not
|
||||
// the URL or the command they produce, and not a target written in the
|
||||
// template itself. These functions run the gate inside the tool, on what the
|
||||
// tool is about to run, before anything is sent or started.
|
||||
|
||||
// renderedPasses bounds the request-and-recheck loop of a rendered target, as
|
||||
// the hook bounds its own: a request that was decided in between, or addresses
|
||||
// that changed again after an approval, send the call back to the pin lookup at
|
||||
// most this many times.
|
||||
const renderedPasses = 3
|
||||
|
||||
// GateRenderedURL is the target gate for the URL a custom http tool renders. It
|
||||
// reads the target as the HTTP client will: the host of the parsed URL and its
|
||||
// port, else the default port of the scheme. A URL whose target cannot be read
|
||||
// is blocked with TARGET_GATE_UNAVAILABLE: the call is not executed and no
|
||||
// approval is requested.
|
||||
//
|
||||
// The result is empty when the tool may run. With the gate off, or with no
|
||||
// scope in ctx (the operator's test endpoint), it returns at once: nothing is
|
||||
// parsed, resolved, read or written.
|
||||
func GateRenderedURL(ctx context.Context, toolName, rawURL string) TargetResult {
|
||||
scope, on := renderedScope(ctx)
|
||||
if !on {
|
||||
return TargetResult{}
|
||||
}
|
||||
input := renderedInput("url", rawURL)
|
||||
t, err := urlTarget(rawURL)
|
||||
if err != nil {
|
||||
return undeterminedTarget(ctx, scope, toolName, input, "the target of the rendered URL could not be read ("+err.Error()+")")
|
||||
}
|
||||
return gateRendered(ctx, scope, toolName, input, []Target{t})
|
||||
}
|
||||
|
||||
// GateRenderedCommand is the target gate for the command a custom command tool
|
||||
// renders, read with the command-string rules of Bash. readings are other
|
||||
// renderings of the same command that are read as well: the shell quoting of a
|
||||
// parameter placed inside a host name hides that name from the rules, and the
|
||||
// shell removes the quotes before the command runs. A command that names no
|
||||
// target passes, as it does in Bash: what a command hides is advisory (D9).
|
||||
//
|
||||
// The result is empty when the tool may run. With the gate off, or with no
|
||||
// scope in ctx, it returns at once.
|
||||
func GateRenderedCommand(ctx context.Context, toolName, command string, readings ...string) TargetResult {
|
||||
scope, on := renderedScope(ctx)
|
||||
if !on {
|
||||
return TargetResult{}
|
||||
}
|
||||
var l targetList
|
||||
l.scan(command)
|
||||
for _, r := range readings {
|
||||
l.scan(r)
|
||||
}
|
||||
return gateRendered(ctx, scope, toolName, renderedInput("command", command), l.targets)
|
||||
}
|
||||
|
||||
func renderedScope(ctx context.Context) (TargetScope, bool) {
|
||||
scope, ok := TargetScopeFromContext(ctx)
|
||||
return scope, ok && scope.Enabled != nil && scope.Enabled()
|
||||
}
|
||||
|
||||
// renderedInput is the tool input recorded with a rendered target: what the
|
||||
// tool runs, which is what the human reads on the approval.
|
||||
func renderedInput(field, value string) []byte {
|
||||
b, _ := json.Marshal(map[string]string{field: value})
|
||||
return b
|
||||
}
|
||||
|
||||
// gateRendered runs the gate decision on the targets of a rendered call, asks
|
||||
// for the approval of those no pin covers, and checks again after the
|
||||
// decisions. It never merges an ask: the rules and the judge decided at the
|
||||
// hook, before the tool ran.
|
||||
func gateRendered(ctx context.Context, scope TargetScope, toolName string, input []byte, targets []Target) TargetResult {
|
||||
if len(targets) == 0 {
|
||||
return TargetResult{}
|
||||
}
|
||||
i := scope.Interceptor
|
||||
if i == nil {
|
||||
return TargetResult{Code: TargetGateUnavailable, Host: targets[0].Host, Port: targets[0].Port,
|
||||
Reason: "the session's target scope carries no gate"}
|
||||
}
|
||||
ctx = WithCall(ctx, toolName, input)
|
||||
check := func() ([]TargetRequest, TargetResult) {
|
||||
if !scope.Enabled() {
|
||||
return nil, TargetResult{}
|
||||
}
|
||||
c := &targetCheck{i: i, ctx: ctx, scope: scope, tool: toolName, input: input}
|
||||
return c.run(targets)
|
||||
}
|
||||
unpinned, res := check()
|
||||
for pass := 0; res.Code == "" && len(unpinned) > 0; pass++ {
|
||||
if pass == renderedPasses {
|
||||
t := unpinned[0]
|
||||
return TargetResult{Code: TargetApprovalPending, Host: t.Host, Port: t.Port,
|
||||
Reason: fmt.Sprintf("the target is still not covered by a pin after %d approval rounds; its addresses changed again", renderedPasses)}
|
||||
}
|
||||
if res = i.RequestTargets(ctx, scope, toolName, input, unpinned, nil); res.Code != "" {
|
||||
return res
|
||||
}
|
||||
unpinned, res = check()
|
||||
}
|
||||
return res
|
||||
}
|
||||
|
||||
// undeterminedTarget blocks a call whose target cannot be read, and logs the
|
||||
// block as a decided row.
|
||||
func undeterminedTarget(ctx context.Context, scope TargetScope, toolName string, input []byte, reason string) TargetResult {
|
||||
res := TargetResult{Code: TargetGateUnavailable, Reason: reason}
|
||||
if scope.Interceptor == nil {
|
||||
return res
|
||||
}
|
||||
ref, err := scope.ref()
|
||||
if err != nil {
|
||||
return res
|
||||
}
|
||||
c := &targetCheck{i: scope.Interceptor, ctx: WithCall(ctx, toolName, input), scope: scope, ref: ref, tool: toolName, input: input}
|
||||
_, res = c.block(TargetGateUnavailable, "gate_unavailable", &checkedTarget{}, reason)
|
||||
return res
|
||||
}
|
||||
|
||||
// urlTarget reads the target of a URL the way net/http dials it: the host of
|
||||
// the parsed URL, and its port, else the default port of the scheme.
|
||||
func urlTarget(raw string) (Target, error) {
|
||||
u, err := url.Parse(raw)
|
||||
if err != nil {
|
||||
return Target{}, err
|
||||
}
|
||||
name := u.Hostname()
|
||||
if name == "" {
|
||||
return Target{}, errors.New("no host")
|
||||
}
|
||||
kind, host, ok := normalizeHost(name, true)
|
||||
if !ok {
|
||||
return Target{}, fmt.Errorf("host %q", name)
|
||||
}
|
||||
t := Target{Kind: kind, Host: host, Port: schemePorts[strings.ToLower(u.Scheme)]}
|
||||
if p := u.Port(); p != "" {
|
||||
if t.Port = validPort(p); t.Port == 0 {
|
||||
return Target{}, fmt.Errorf("port %q", p)
|
||||
}
|
||||
}
|
||||
return t, nil
|
||||
}
|
||||
@@ -0,0 +1,177 @@
|
||||
package intercept
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// The gate inside the custom tools: GateRenderedURL and GateRenderedCommand
|
||||
// read what the tool is about to run. Without a database: the cases that reach
|
||||
// a request are covered by the server tests.
|
||||
|
||||
func (g *gateFixture) ctx() context.Context {
|
||||
return WithTargetScope(context.Background(), g.scope)
|
||||
}
|
||||
|
||||
// With the gate off, or without a scope (the operator's test endpoint), the
|
||||
// rendered target is not even parsed: nothing is read, resolved or logged, and
|
||||
// a URL the gate could not read passes too, as it did before the gate.
|
||||
func TestRenderedGateOffTouchesNothing(t *testing.T) {
|
||||
g := newGate(t)
|
||||
g.res.set("app.example.com", "203.0.113.7")
|
||||
for name, ctx := range map[string]context.Context{
|
||||
"no scope": context.Background(),
|
||||
"no flag": WithTargetScope(context.Background(), TargetScope{TaskID: "7", WaitBound: time.Minute, Interceptor: g.ic}),
|
||||
"flag off": WithTargetScope(context.Background(), TargetScope{TaskID: "7", Enabled: func() bool { return false }, Interceptor: g.ic}),
|
||||
} {
|
||||
for _, res := range []TargetResult{
|
||||
GateRenderedURL(ctx, "probe", appURL),
|
||||
GateRenderedURL(ctx, "probe", "http://a b/"),
|
||||
GateRenderedCommand(ctx, "probe", "curl "+appURL),
|
||||
} {
|
||||
if res != (TargetResult{}) {
|
||||
t.Errorf("%s: %+v, want nothing", name, res)
|
||||
}
|
||||
}
|
||||
}
|
||||
if g.store.reads != 0 || len(g.store.logged) != 0 || g.res.total() != 0 {
|
||||
t.Errorf("gate off: %d store reads, %d rows, %d lookups, want none", g.store.reads, len(g.store.logged), g.res.total())
|
||||
}
|
||||
}
|
||||
|
||||
// The target of a rendered URL is read as net/http dials it: the host of the
|
||||
// parsed URL, the explicit port, else the default of the scheme.
|
||||
func TestRenderedURLTarget(t *testing.T) {
|
||||
for _, c := range []struct {
|
||||
raw string
|
||||
want Target
|
||||
}{
|
||||
{"https://App.Example.COM./x?next=http://other.example/", Target{TargetDomain, "app.example.com", 443}},
|
||||
{"http://user:pw@203.0.113.7:8080/a", Target{TargetIP, "203.0.113.7", 8080}},
|
||||
{"http://[::ffff:203.0.113.7]/", Target{TargetIP, "203.0.113.7", 80}},
|
||||
{"http://[2001:db8::1]:8443/", Target{TargetIP, "2001:db8::1", 8443}},
|
||||
{"http://2130706433/", Target{TargetIP, "127.0.0.1", 80}},
|
||||
{"HTTPS://bücher.example/", Target{TargetDomain, "xn--bcher-kva.example", 443}},
|
||||
// The client's reading of the authority, not the shell's: a.example; is userinfo.
|
||||
{"http://a.example;@b.example/", Target{TargetDomain, "b.example", 80}},
|
||||
} {
|
||||
got, err := urlTarget(c.raw)
|
||||
if err != nil || got != c.want {
|
||||
t.Errorf("%q: %+v, %v; want %+v", c.raw, got, err, c.want)
|
||||
}
|
||||
}
|
||||
for _, raw := range []string{
|
||||
"http://a b/", // not a URL the client accepts
|
||||
`http://a.example\@b.example/`, // idem
|
||||
"http://[::1/", // idem
|
||||
"http:///path", // no host
|
||||
"app.example.com/path", // no scheme: no host
|
||||
"http://app.example.com:99999/x", // no valid port
|
||||
} {
|
||||
if got, err := urlTarget(raw); err == nil {
|
||||
t.Errorf("%q: %+v, want no target", raw, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// A rendered URL whose target cannot be read is blocked: no lookup, no
|
||||
// request, one decided row that records the URL.
|
||||
func TestRenderedURLWithoutTargetFailsClosed(t *testing.T) {
|
||||
g := newGate(t)
|
||||
for n, raw := range []string{"http://a b/", "http:///path", "app.example.com/path", "http://app.example.com:99999/"} {
|
||||
res := GateRenderedURL(g.ctx(), "probe_http", raw)
|
||||
if res.Code != TargetGateUnavailable || !strings.Contains(res.Reason, "could not be read") {
|
||||
t.Fatalf("%q: %+v, want %s", raw, res, TargetGateUnavailable)
|
||||
}
|
||||
if len(g.store.logged) != n+1 {
|
||||
t.Fatalf("%q: %d rows, want %d", raw, len(g.store.logged), n+1)
|
||||
}
|
||||
row := g.store.logged[n]
|
||||
var input map[string]string
|
||||
_ = json.Unmarshal(row.ToolInput, &input)
|
||||
if row.Status != "denied" || row.ToolName != "probe_http" || input["url"] != raw || row.Audit.Targets[0].Outcome != "gate_unavailable" {
|
||||
t.Fatalf("%q: row = %+v (input %s)", raw, row, row.ToolInput)
|
||||
}
|
||||
}
|
||||
if g.res.total() != 0 {
|
||||
t.Fatalf("%d lookups, want none", g.res.total())
|
||||
}
|
||||
}
|
||||
|
||||
// A gated scope without the interceptor that owns the gate cannot decide: the
|
||||
// call is blocked, never let through.
|
||||
func TestRenderedGateWithoutInterceptorFailsClosed(t *testing.T) {
|
||||
ctx := WithTargetScope(context.Background(), TargetScope{TaskID: "7", Enabled: func() bool { return true }})
|
||||
for name, res := range map[string]TargetResult{
|
||||
"url": GateRenderedURL(ctx, "probe", appURL),
|
||||
"unreadable url": GateRenderedURL(ctx, "probe", "http://a b/"),
|
||||
"command": GateRenderedCommand(ctx, "probe", "curl "+appURL),
|
||||
} {
|
||||
if res.Code != TargetGateUnavailable {
|
||||
t.Errorf("%s: %+v, want %s", name, res, TargetGateUnavailable)
|
||||
}
|
||||
}
|
||||
// A command that names no target needs no gate.
|
||||
if res := GateRenderedCommand(ctx, "probe", "ls -la"); res != (TargetResult{}) {
|
||||
t.Errorf("no target: %+v, want nothing", res)
|
||||
}
|
||||
}
|
||||
|
||||
// The rendered target goes through the same decision as a target at the hook:
|
||||
// block rules, resolution, pins. A pinned target passes, a blocked one is
|
||||
// blocked without a request.
|
||||
func TestRenderedTargetUsesTheGateDecision(t *testing.T) {
|
||||
g := newGate(t)
|
||||
g.res.set("app.example.com", "203.0.113.7")
|
||||
g.store.approve(TargetRequest{Host: "app.example.com", Port: 443, ResolvedIPs: []string{"203.0.113.7"}})
|
||||
g.store.approve(TargetRequest{Host: "198.51.100.9", Port: 8080, ResolvedIPs: []string{"198.51.100.9"}})
|
||||
|
||||
for name, res := range map[string]TargetResult{
|
||||
"url": GateRenderedURL(g.ctx(), "probe", appURL),
|
||||
"url, ip and port": GateRenderedURL(g.ctx(), "probe", "http://198.51.100.9:8080/x"),
|
||||
"command": GateRenderedCommand(g.ctx(), "probe", "curl -s '"+appURL+"'"),
|
||||
"command, no target": GateRenderedCommand(g.ctx(), "probe", "nmap -sV app.example.com"),
|
||||
"command, ip literal": GateRenderedCommand(g.ctx(), "probe", "nc 198.51.100.9:8080"),
|
||||
} {
|
||||
if res != (TargetResult{}) {
|
||||
t.Errorf("%s: %+v, want the call to pass", name, res)
|
||||
}
|
||||
}
|
||||
|
||||
g.res.errs["down.example"] = context.DeadlineExceeded
|
||||
if res := GateRenderedURL(g.ctx(), "probe", "https://down.example/"); res.Code != TargetResolveFailed || res.Host != "down.example" {
|
||||
t.Fatalf("unresolvable: %+v", res)
|
||||
}
|
||||
g.store.rules = []db.AssetInterceptRule{{Enabled: true, Kind: "exact_domain", Pattern: "app.example.com"}}
|
||||
for name, res := range map[string]TargetResult{
|
||||
"url": GateRenderedURL(g.ctx(), "probe", appURL),
|
||||
"command": GateRenderedCommand(g.ctx(), "probe", "curl "+appURL),
|
||||
} {
|
||||
if res.Code != TargetRuleBlocked || res.Host != "app.example.com" || res.Port != 443 {
|
||||
t.Errorf("%s: %+v, want %s on the pinned host", name, res, TargetRuleBlocked)
|
||||
}
|
||||
}
|
||||
if !strings.Contains(strings.Join(g.store.outcomes(), " "), "denied/rule_blocked") {
|
||||
t.Fatalf("rows = %v, want the block logged", g.store.outcomes())
|
||||
}
|
||||
}
|
||||
|
||||
// The shell quoting of a parameter inside a host name hides the name from the
|
||||
// command-string rules; the reading without the quoting is read as well.
|
||||
func TestRenderedCommandReadsTheUnquotedRendering(t *testing.T) {
|
||||
g := newGate(t)
|
||||
g.store.rules = []db.AssetInterceptRule{{Enabled: true, Kind: "exact_domain", Pattern: "api.example.com"}}
|
||||
quoted := "curl http://'api'.example.com:8080/"
|
||||
if res := GateRenderedCommand(g.ctx(), "probe", quoted); res != (TargetResult{}) {
|
||||
t.Fatalf("quoted only: %+v; the rules were expected to miss the name", res)
|
||||
}
|
||||
res := GateRenderedCommand(g.ctx(), "probe", quoted, "curl http://api.example.com:8080/")
|
||||
if res.Code != TargetRuleBlocked || res.Host != "api.example.com" || res.Port != 8080 {
|
||||
t.Fatalf("with the unquoted reading: %+v, want %s on api.example.com:8080", res, TargetRuleBlocked)
|
||||
}
|
||||
}
|
||||
+29
-12
@@ -620,9 +620,14 @@ func (s *Server) StartTriggeredRun(agentKey, title, message string, taskID int64
|
||||
return
|
||||
}
|
||||
cfg := s.readTriggerBehavior(agentKey) // DB read BEFORE the lock (never under queueMu)
|
||||
// A conversation triggered by a gated task is gated (the flag only, never pins).
|
||||
gated := false
|
||||
if t, ok := s.m.Task(strconv.FormatInt(taskID, 10)); ok && taskID > 0 {
|
||||
gated = t.targetGate.Load()
|
||||
}
|
||||
s.queueMu.Lock()
|
||||
s.triggerCfg[agentKey] = cfg
|
||||
s.triggerQ[agentKey] = append(s.triggerQ[agentKey], triggeredRun{agentKey: agentKey, title: title, message: message, taskID: taskID, taskDesc: taskDesc, taskGoal: taskGoal, mergeable: mergeable})
|
||||
s.triggerQ[agentKey] = append(s.triggerQ[agentKey], triggeredRun{agentKey: agentKey, title: title, message: message, taskID: taskID, taskDesc: taskDesc, taskGoal: taskGoal, mergeable: mergeable, targetGate: gated})
|
||||
s.pumpLocked(agentKey)
|
||||
s.queueMu.Unlock()
|
||||
}
|
||||
@@ -749,15 +754,26 @@ func mergeTriggeredRuns(items []triggeredRun) triggeredRun {
|
||||
fmt.Fprintf(&b, "\n── trigger %d ──\n%s\n", i+1, it.message)
|
||||
}
|
||||
return triggeredRun{
|
||||
agentKey: first.agentKey,
|
||||
title: fmt.Sprintf("Merged triggers · task#%d · %d items", first.taskID, len(items)),
|
||||
message: b.String(),
|
||||
taskID: first.taskID,
|
||||
mergeable: true,
|
||||
agentKey: first.agentKey,
|
||||
title: fmt.Sprintf("Merged triggers · task#%d · %d items", first.taskID, len(items)),
|
||||
message: b.String(),
|
||||
taskID: first.taskID,
|
||||
mergeable: true,
|
||||
targetGate: anyTriggerGated(items),
|
||||
// taskDesc/taskGoal left empty: header already embedded above.
|
||||
}
|
||||
}
|
||||
|
||||
// anyTriggerGated reports whether any fire of a merged run came from a gated task.
|
||||
func anyTriggerGated(items []triggeredRun) bool {
|
||||
for _, it := range items {
|
||||
if it.targetGate {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// mergeAllRuns folds the ENTIRE queued burst into one run (merge_mode='all'),
|
||||
// regardless of type. Events are grouped by task (tasks in first-appearance order,
|
||||
// events in their original order within a task), so each task's context header — and
|
||||
@@ -791,11 +807,12 @@ func mergeAllRuns(items []triggeredRun) triggeredRun {
|
||||
}
|
||||
}
|
||||
return triggeredRun{
|
||||
agentKey: first.agentKey,
|
||||
title: fmt.Sprintf("Merged triggers · all · %d items", len(items)),
|
||||
message: b.String(),
|
||||
taskID: first.taskID,
|
||||
mergeable: true,
|
||||
agentKey: first.agentKey,
|
||||
title: fmt.Sprintf("Merged triggers · all · %d items", len(items)),
|
||||
message: b.String(),
|
||||
taskID: first.taskID,
|
||||
mergeable: true,
|
||||
targetGate: anyTriggerGated(items),
|
||||
// taskDesc/taskGoal left empty: per-task headers already embedded above.
|
||||
}
|
||||
}
|
||||
@@ -810,7 +827,7 @@ func (s *Server) runTriggeredRun(item triggeredRun) {
|
||||
}
|
||||
}()
|
||||
pg := s.m.pg
|
||||
c, err := pg.CreateConversation(item.agentKey, firstLine(item.title, 60), nil)
|
||||
c, err := pg.CreateConversationGated(item.agentKey, firstLine(item.title, 60), nil, item.targetGate)
|
||||
if err != nil {
|
||||
log.Printf("[trigger] create conversation for %s failed: %v", item.agentKey, err)
|
||||
return
|
||||
|
||||
+19
-6
@@ -18,6 +18,8 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/guard"
|
||||
"github.com/Autumn-27/artex/intercept"
|
||||
"github.com/Autumn-27/norma/permission"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
)
|
||||
@@ -163,11 +165,11 @@ func (s *Server) pgTestCustomTool(w http.ResponseWriter, r *http.Request) {
|
||||
var res actool.Result
|
||||
switch req.Kind {
|
||||
case "command":
|
||||
res, _ = s.runCommandTool(ctx, req.Exec, params, tc)
|
||||
res, _ = s.runCommandTool(ctx, "test", req.Exec, params, tc)
|
||||
case "script":
|
||||
res, _ = s.runScriptTool(ctx, "test", req.Exec, params, tc)
|
||||
case "http":
|
||||
res, _ = s.runHTTPTool(ctx, req.Exec, params, tc)
|
||||
res, _ = s.runHTTPTool(ctx, "test", req.Exec, params, tc)
|
||||
case "shell":
|
||||
writeErr(w, 400, "shell-kind tools are bash environment declarations with no executable content")
|
||||
return
|
||||
@@ -275,11 +277,11 @@ func (s *Server) buildCustomTool(t *db.Tool) actool.CoreTool {
|
||||
}
|
||||
switch kind {
|
||||
case "command":
|
||||
return s.runCommandTool(ctx, execRaw, params, tc)
|
||||
return s.runCommandTool(ctx, key, execRaw, params, tc)
|
||||
case "script":
|
||||
return s.runScriptTool(ctx, key, execRaw, params, tc)
|
||||
case "http":
|
||||
return s.runHTTPTool(ctx, execRaw, params, tc)
|
||||
return s.runHTTPTool(ctx, key, execRaw, params, tc)
|
||||
default:
|
||||
return actool.Errorf("unknown custom tool kind: " + kind), nil
|
||||
}
|
||||
@@ -328,13 +330,18 @@ func ensureSchema(raw json.RawMessage) map[string]any {
|
||||
|
||||
// ---------- command: render the command → reuse the Bash low-level run ----------
|
||||
|
||||
func (s *Server) runCommandTool(ctx context.Context, execRaw json.RawMessage, params map[string]any, tc *actool.ToolContext) (actool.Result, error) {
|
||||
func (s *Server) runCommandTool(ctx context.Context, key string, execRaw json.RawMessage, params map[string]any, tc *actool.ToolContext) (actool.Result, error) {
|
||||
var spec commandExec
|
||||
_ = json.Unmarshal(execRaw, &spec)
|
||||
if strings.TrimSpace(spec.Command) == "" {
|
||||
return actool.Errorf("command is empty"), nil
|
||||
}
|
||||
cmd := renderTemplate(spec.Command, params, shellQuote)
|
||||
// The hook saw the parameters, not the command they render: the target gate
|
||||
// reads the rendered command here, also without the quoting of the parameters.
|
||||
if res := intercept.GateRenderedCommand(ctx, key, cmd, renderTemplate(spec.Command, params, identity)); res.Code != "" {
|
||||
return actool.Errorf(guard.TargetBlockMessage(res)), nil
|
||||
}
|
||||
// Reuse the low-level run Bash also uses (via the Bash CoreTool.Call): it
|
||||
// automatically inherits the safety floor/timeout/proxy env/output overflow. The tool is a peer of Bash sharing that low level, not routed through the Bash tool for the model to call.
|
||||
bashIn, _ := json.Marshal(map[string]any{"command": cmd})
|
||||
@@ -417,7 +424,7 @@ func execPython(ctx context.Context, interp, key, code string, params map[string
|
||||
|
||||
// ---------- http: native request + proxy ----------
|
||||
|
||||
func (s *Server) runHTTPTool(ctx context.Context, execRaw json.RawMessage, params map[string]any, tc *actool.ToolContext) (actool.Result, error) {
|
||||
func (s *Server) runHTTPTool(ctx context.Context, key string, execRaw json.RawMessage, params map[string]any, tc *actool.ToolContext) (actool.Result, error) {
|
||||
var spec httpExec
|
||||
_ = json.Unmarshal(execRaw, &spec)
|
||||
method := strings.ToUpper(strings.TrimSpace(spec.Method))
|
||||
@@ -428,6 +435,12 @@ func (s *Server) runHTTPTool(ctx context.Context, execRaw json.RawMessage, param
|
||||
if strings.TrimSpace(rawURL) == "" {
|
||||
return actool.Errorf("http url is empty"), nil
|
||||
}
|
||||
// The hook saw the parameters, not the URL they render (nor a URL written in
|
||||
// the template): the target gate reads the rendered URL here, before the
|
||||
// tool's own timeout starts, so a wait for an approval does not consume it.
|
||||
if res := intercept.GateRenderedURL(ctx, key, rawURL); res.Code != "" {
|
||||
return actool.Errorf(guard.TargetBlockMessage(res)), nil
|
||||
}
|
||||
var bodyReader io.Reader
|
||||
if spec.Body != "" {
|
||||
bodyReader = strings.NewReader(renderTemplate(spec.Body, params, identity))
|
||||
|
||||
@@ -0,0 +1,619 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"iter"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/netip"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/intercept"
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
)
|
||||
|
||||
// Step 1.8: the targets the custom tools render. The hook reads the parameters
|
||||
// of a custom tool, not the URL or the command they produce: the tool runs the
|
||||
// gate on what it is about to run. Plus the end-to-end run of a gated worker
|
||||
// and the run of a task with the gate off.
|
||||
|
||||
// countingServer is an HTTP server on 127.0.0.1 that counts the requests it
|
||||
// receives. Its host:port is an IP literal: the gate never resolves it.
|
||||
func countingServer(t *testing.T) (srv *httptest.Server, hits *atomic.Int64, host, port string) {
|
||||
t.Helper()
|
||||
hits = &atomic.Int64{}
|
||||
srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
hits.Add(1)
|
||||
w.Header().Set("Content-Type", "text/plain")
|
||||
_, _ = w.Write([]byte("probe-reached"))
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
u, err := url.Parse(srv.URL)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
host, port, err = net.SplitHostPort(u.Host)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return srv, hits, host, port
|
||||
}
|
||||
|
||||
// gatedCtx is the context of a worker tool call of task: its approval scope
|
||||
// with the given wait bound.
|
||||
func gatedCtx(task *Task, wait time.Duration) context.Context {
|
||||
ctx := intercept.WithTaskContext(context.Background(), task.ID, "worker", nil)
|
||||
return intercept.WithTargetScope(ctx, task.targetScope(wait))
|
||||
}
|
||||
|
||||
func execJSON(t *testing.T, v any) json.RawMessage {
|
||||
t.Helper()
|
||||
b, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
// callTool runs a custom tool the way the harness does and returns its text.
|
||||
func callTool(t *testing.T, tool actool.CoreTool, ctx context.Context, input string, tc *actool.ToolContext) (string, bool) {
|
||||
t.Helper()
|
||||
res, err := tool.Call(ctx, json.RawMessage(input), tc)
|
||||
if err != nil {
|
||||
t.Fatalf("%s: %v", tool.Name(), err)
|
||||
}
|
||||
return res.Flatten(), res.IsError
|
||||
}
|
||||
|
||||
// openTargetRequest returns the open target request of the task on key.
|
||||
func openTargetRequest(t *testing.T, pg *db.DB, taskID, key string) (id int64, toolName, toolInput string, ok bool) {
|
||||
t.Helper()
|
||||
err := pg.QueryRow(`SELECT id, tool_name, tool_input::text FROM intercept_pending
|
||||
WHERE task_id=$1 AND status='pending' AND target_key=$2`, taskID, key).Scan(&id, &toolName, &toolInput)
|
||||
if err != nil {
|
||||
return 0, "", "", false
|
||||
}
|
||||
return id, toolName, toolInput, true
|
||||
}
|
||||
|
||||
func targetRowCount(t *testing.T, pg *db.DB, taskID string) int {
|
||||
t.Helper()
|
||||
var n int
|
||||
if err := pg.QueryRow(`SELECT count(*) FROM intercept_pending WHERE task_id=$1 AND (decision_source='target' OR target_key IS NOT NULL)`, taskID).Scan(&n); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
// A custom http tool is gated on the URL it renders, whether the target comes
|
||||
// from a parameter (here split over two parameters, so the hook reads no
|
||||
// target in them) or is written in the template. Nothing is sent before an
|
||||
// approval; an approval lets both through; a revoked pin asks again.
|
||||
func TestCustomHTTPRenderedTargetGated(t *testing.T) {
|
||||
g := newGateServer(t)
|
||||
on := true
|
||||
task := g.newTask(t, db.TaskCreateOptions{TargetGate: &on})
|
||||
pg := g.m.pg
|
||||
srv, hits, host, port := countingServer(t)
|
||||
key := net.JoinHostPort(host, port)
|
||||
dot := strings.LastIndexByte(host, '.')
|
||||
|
||||
rendered := g.s.buildCustomTool(&db.Tool{Key: "probe_http", Kind: "http",
|
||||
Schema: json.RawMessage(`{"type":"object","properties":{"last":{"type":"string"},"port":{"type":"string"}}}`),
|
||||
Exec: execJSON(t, map[string]any{"method": "GET", "url": "http://" + host[:dot+1] + "{last}:{port}/rendered"})})
|
||||
params := fmt.Sprintf(`{"last":%q,"port":%q}`, host[dot+1:], port)
|
||||
fixed := g.s.buildCustomTool(&db.Tool{Key: "probe_fixed", Kind: "http",
|
||||
Exec: execJSON(t, map[string]any{"method": "GET", "url": srv.URL + "/fixed"})})
|
||||
// The hook cannot see either target: neither input names one.
|
||||
for _, in := range []struct{ tool, input string }{{"probe_http", params}, {"probe_fixed", `{}`}} {
|
||||
if got := intercept.ExtractTargets(in.tool, []byte(in.input)); len(got) != 0 {
|
||||
t.Fatalf("%s: the hook reads %+v; the test needs a target only the rendering shows", in.tool, got)
|
||||
}
|
||||
}
|
||||
|
||||
// The template's own URL, wait bound 0: blocked at once, nothing sent, the
|
||||
// request open with the rendered URL as its input.
|
||||
out, isErr := callTool(t, fixed, gatedCtx(task, 0), `{}`, nil)
|
||||
if !isErr || !strings.Contains(out, "TARGET_APPROVAL_PENDING") || !strings.Contains(out, key) || hits.Load() != 0 {
|
||||
t.Fatalf("hard-coded URL: error=%v out=%q hits=%d", isErr, out, hits.Load())
|
||||
}
|
||||
id, toolName, toolInput, ok := openTargetRequest(t, pg, task.ID, key)
|
||||
if !ok || toolName != "probe_fixed" || !strings.Contains(toolInput, srv.URL+"/fixed") {
|
||||
t.Fatalf("request row: ok=%v tool=%q input=%s", ok, toolName, toolInput)
|
||||
}
|
||||
|
||||
// The rendered URL names the same host and port: the call attaches to the
|
||||
// open request and waits. One click lets it through.
|
||||
type result struct {
|
||||
out string
|
||||
isErr bool
|
||||
}
|
||||
done := make(chan result, 1)
|
||||
go func() {
|
||||
res, err := rendered.Call(gatedCtx(task, 30*time.Second), json.RawMessage(params), nil)
|
||||
if err != nil {
|
||||
done <- result{err.Error(), true}
|
||||
return
|
||||
}
|
||||
done <- result{res.Flatten(), res.IsError}
|
||||
}()
|
||||
time.Sleep(300 * time.Millisecond)
|
||||
if hits.Load() != 0 {
|
||||
t.Fatal("the rendered URL was fetched before the approval")
|
||||
}
|
||||
if n := targetRowCount(t, pg, task.ID); n != 1 {
|
||||
t.Fatalf("%d target rows, want the one request", n)
|
||||
}
|
||||
if code, body := g.do(t, "POST", fmt.Sprintf("/api/intercept/pending/%d/decide", id), `{"decision":"allowed"}`); code != 200 {
|
||||
t.Fatalf("decide: %d %v", code, body)
|
||||
}
|
||||
select {
|
||||
case r := <-done:
|
||||
if r.isErr || !strings.Contains(r.out, "probe-reached") || hits.Load() != 1 {
|
||||
t.Fatalf("after the approval: error=%v out=%q hits=%d", r.isErr, r.out, hits.Load())
|
||||
}
|
||||
case <-time.After(30 * time.Second):
|
||||
t.Fatal("the waiting call did not end after the approval")
|
||||
}
|
||||
if out, isErr := callTool(t, fixed, gatedCtx(task, 0), `{}`, nil); isErr || hits.Load() != 2 {
|
||||
t.Fatalf("pinned, hard-coded URL: error=%v out=%q hits=%d", isErr, out, hits.Load())
|
||||
}
|
||||
|
||||
// Revoked: the next call asks again and sends nothing.
|
||||
pins, err := pg.ListTargetPins(taskRef(t, task))
|
||||
if err != nil || len(pins) != 1 || pins[0].Host != host || fmt.Sprint(pins[0].Port) != port {
|
||||
t.Fatalf("pins = %+v, %v", pins, err)
|
||||
}
|
||||
if code, _ := g.do(t, "POST", fmt.Sprintf("/api/tasks/%s/target-pins/%d/revoke", task.ID, pins[0].ID), ""); code != 200 {
|
||||
t.Fatalf("revoke: %d", code)
|
||||
}
|
||||
out, isErr = callTool(t, rendered, gatedCtx(task, 0), params, nil)
|
||||
if !isErr || !strings.Contains(out, "TARGET_APPROVAL_PENDING") || hits.Load() != 2 {
|
||||
t.Fatalf("revoked: error=%v out=%q hits=%d", isErr, out, hits.Load())
|
||||
}
|
||||
if again, toolName, _, ok := openTargetRequest(t, pg, task.ID, key); !ok || again == id || toolName != "probe_http" {
|
||||
t.Fatalf("after the revoke: open request %d (ok=%v, tool %q), want a new one", again, ok, toolName)
|
||||
}
|
||||
|
||||
// A URL whose target cannot be read is blocked, never sent.
|
||||
broken := g.s.buildCustomTool(&db.Tool{Key: "probe_broken", Kind: "http",
|
||||
Schema: json.RawMessage(`{"type":"object","properties":{"host":{"type":"string"}}}`),
|
||||
Exec: execJSON(t, map[string]any{"url": "http://{host}/x"})})
|
||||
out, isErr = callTool(t, broken, gatedCtx(task, 0), `{"host":"a b"}`, nil)
|
||||
if !isErr || !strings.Contains(out, "TARGET_GATE_UNAVAILABLE") {
|
||||
t.Fatalf("unreadable URL: error=%v out=%q", isErr, out)
|
||||
}
|
||||
}
|
||||
|
||||
// A custom command tool is gated on the command it renders, with a target
|
||||
// written in the template or formed by a parameter inside a host name. The
|
||||
// command does not run before an approval.
|
||||
func TestCustomCommandRenderedTargetGated(t *testing.T) {
|
||||
g := newGateServer(t)
|
||||
on := true
|
||||
task := g.newTask(t, db.TaskCreateOptions{TargetGate: &on})
|
||||
pg := g.m.pg
|
||||
res := &countingResolver{answers: map[string][]string{"api.example.com": {"198.51.100.7"}}}
|
||||
g.m.interceptor.SetResolver(res)
|
||||
dir := t.TempDir()
|
||||
marker := filepath.Join(dir, "marker")
|
||||
ran := func() bool {
|
||||
_, err := os.Stat(marker)
|
||||
return err == nil
|
||||
}
|
||||
|
||||
fixed := g.s.buildCustomTool(&db.Tool{Key: "probe_cmd_fixed", Kind: "command",
|
||||
Exec: execJSON(t, map[string]any{"command": "echo reached > '" + marker + "' && : http://198.51.100.7:8080/"})})
|
||||
glued := g.s.buildCustomTool(&db.Tool{Key: "probe_cmd", Kind: "command",
|
||||
Schema: json.RawMessage(`{"type":"object","properties":{"sub":{"type":"string"}}}`),
|
||||
Exec: execJSON(t, map[string]any{"command": "echo reached > '" + marker + "' && : http://{sub}.example.com:8080/"})})
|
||||
if got := intercept.ExtractTargets("probe_cmd", []byte(`{"sub":"api"}`)); len(got) != 0 {
|
||||
t.Fatalf("the hook reads %+v; the test needs a target only the rendering shows", got)
|
||||
}
|
||||
tc := &actool.ToolContext{WorkingDir: dir}
|
||||
|
||||
for _, c := range []struct {
|
||||
name string
|
||||
tool actool.CoreTool
|
||||
input string
|
||||
key string
|
||||
}{
|
||||
{"template", fixed, `{}`, "198.51.100.7:8080"},
|
||||
{"parameter inside the host", glued, `{"sub":"api"}`, "api.example.com:8080"},
|
||||
} {
|
||||
out, isErr := callTool(t, c.tool, gatedCtx(task, 0), c.input, tc)
|
||||
if !isErr || !strings.Contains(out, "TARGET_APPROVAL_PENDING") || ran() {
|
||||
t.Fatalf("%s: error=%v out=%q ran=%v, want a block before the command runs", c.name, isErr, out, ran())
|
||||
}
|
||||
id, toolName, _, ok := openTargetRequest(t, pg, task.ID, c.key)
|
||||
if !ok || toolName != c.tool.Name() {
|
||||
t.Fatalf("%s: no open request on %s for %s", c.name, c.key, c.tool.Name())
|
||||
}
|
||||
if code, body := g.do(t, "POST", fmt.Sprintf("/api/intercept/pending/%d/decide", id), `{"decision":"allowed"}`); code != 200 {
|
||||
t.Fatalf("%s: decide: %d %v", c.name, code, body)
|
||||
}
|
||||
if out, isErr := callTool(t, c.tool, gatedCtx(task, 0), c.input, tc); isErr || !ran() {
|
||||
t.Fatalf("%s, approved: error=%v out=%q ran=%v", c.name, isErr, out, ran())
|
||||
}
|
||||
_ = os.Remove(marker)
|
||||
}
|
||||
if res.lookups() == 0 {
|
||||
t.Fatal("the name formed by the parameter was never resolved")
|
||||
}
|
||||
}
|
||||
|
||||
// The custom-tool test endpoint is the operator's own action (D14): it has no
|
||||
// approval scope, so the gate is not on there, whatever the tasks say.
|
||||
func TestCustomToolTestEndpointUngated(t *testing.T) {
|
||||
g := newGateServer(t)
|
||||
on := true
|
||||
g.newTask(t, db.TaskCreateOptions{TargetGate: &on})
|
||||
_, hits, host, port := countingServer(t)
|
||||
before := 0
|
||||
if err := g.m.pg.QueryRow(`SELECT count(*) FROM intercept_pending WHERE decision_source='target' OR target_key IS NOT NULL`).Scan(&before); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
body := fmt.Sprintf(`{"kind":"http","exec":{"url":"http://%s:{port}/test"},"params":{"port":%q}}`, host, port)
|
||||
code, out := g.do(t, "POST", "/api/tools/custom/test", body)
|
||||
if code != 200 || out["is_error"] != false || !strings.Contains(fmt.Sprint(out["output"]), "probe-reached") || hits.Load() != 1 {
|
||||
t.Fatalf("http: %d %v hits=%d", code, out, hits.Load())
|
||||
}
|
||||
code, out = g.do(t, "POST", "/api/tools/custom/test", `{"kind":"command","exec":{"command":": http://198.51.100.7:8080/ && echo ran-ungated"}}`)
|
||||
if code != 200 || !strings.Contains(fmt.Sprint(out["output"]), "ran-ungated") {
|
||||
t.Fatalf("command: %d %v", code, out)
|
||||
}
|
||||
after := 0
|
||||
if err := g.m.pg.QueryRow(`SELECT count(*) FROM intercept_pending WHERE decision_source='target' OR target_key IS NOT NULL`).Scan(&after); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if after != before {
|
||||
t.Fatalf("the test endpoint wrote %d target rows", after-before)
|
||||
}
|
||||
}
|
||||
|
||||
// countingResolver answers from a table and counts every lookup.
|
||||
type countingResolver struct {
|
||||
mu sync.Mutex
|
||||
answers map[string][]string
|
||||
n int
|
||||
}
|
||||
|
||||
func (r *countingResolver) Name() string { return "counting" }
|
||||
|
||||
func (r *countingResolver) LookupIPs(_ context.Context, host string) ([]netip.Addr, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.n++
|
||||
var out []netip.Addr
|
||||
for _, ip := range r.answers[host] {
|
||||
out = append(out, netip.MustParseAddr(ip))
|
||||
}
|
||||
if len(out) == 0 {
|
||||
return nil, fmt.Errorf("no answer for %s", host)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *countingResolver) lookups() int {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return r.n
|
||||
}
|
||||
|
||||
// scriptedCall is one tool call of a toolScript.
|
||||
type scriptedCall struct{ tool, input string }
|
||||
|
||||
func fetches(urls ...string) []scriptedCall {
|
||||
calls := make([]scriptedCall, len(urls))
|
||||
for n, u := range urls {
|
||||
calls[n] = scriptedCall{"WebFetch", fmt.Sprintf(`{"url":%q}`, u)}
|
||||
}
|
||||
return calls
|
||||
}
|
||||
|
||||
// toolScript is a model that makes each of its calls in turn, one per turn,
|
||||
// then ends the run. It hands the result of every call to the test, and waits
|
||||
// before turn n while hold[n] is open.
|
||||
type toolScript struct {
|
||||
calls []scriptedCall
|
||||
hold map[int64]chan struct{}
|
||||
results chan string
|
||||
turn atomic.Int64
|
||||
}
|
||||
|
||||
func (p *toolScript) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
|
||||
n := p.turn.Add(1)
|
||||
if n > 1 {
|
||||
select {
|
||||
case p.results <- lastToolResult(req.Messages):
|
||||
default:
|
||||
}
|
||||
}
|
||||
if h := p.hold[n]; h != nil {
|
||||
select {
|
||||
case <-h:
|
||||
case <-ctx.Done():
|
||||
}
|
||||
}
|
||||
events := []llm.StreamEvent{{Type: llm.SETextDelta, Text: "done"}, {Type: llm.SEMessageDelta, StopReason: "end_turn"}}
|
||||
if int(n) <= len(p.calls) {
|
||||
call := p.calls[n-1]
|
||||
events = []llm.StreamEvent{
|
||||
{Type: llm.SEToolUseStart, ToolID: fmt.Sprintf("call-%d", n), ToolName: call.tool},
|
||||
{Type: llm.SEToolInputJSON, Text: call.input},
|
||||
{Type: llm.SEMessageDelta, StopReason: "tool_use"},
|
||||
}
|
||||
}
|
||||
return func(yield func(llm.StreamEvent, error) bool) {
|
||||
for _, event := range events {
|
||||
if !yield(event, nil) {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (p *toolScript) Complete(ctx context.Context, req llm.CompletionRequest) (llm.Message, string, llm.Usage, error) {
|
||||
return accumulateStreamForTest(ctx, p.Stream, req)
|
||||
}
|
||||
|
||||
func (p *toolScript) next(t *testing.T) string {
|
||||
t.Helper()
|
||||
select {
|
||||
case r := <-p.results:
|
||||
return r
|
||||
case <-time.After(60 * time.Second):
|
||||
t.Fatal("the worker did not hand back a tool result")
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
// lastToolResult is the text of the last tool result the model is shown.
|
||||
func lastToolResult(msgs []llm.Message) string {
|
||||
for i := len(msgs) - 1; i >= 0; i-- {
|
||||
for _, b := range msgs[i].Content {
|
||||
if b.Type != llm.BlockToolResult {
|
||||
continue
|
||||
}
|
||||
var s strings.Builder
|
||||
if b.IsError {
|
||||
s.WriteString("[error] ")
|
||||
}
|
||||
for _, c := range b.Content {
|
||||
s.WriteString(c.Text)
|
||||
}
|
||||
return s.String()
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// runScriptedWorker claims one intent of task and runs it with a worker driven
|
||||
// by p, in the background. extra are tools the worker has besides its own.
|
||||
func runScriptedWorker(t *testing.T, g *gateServer, task *Task, p llm.Provider, extra ...actool.CoreTool) <-chan struct{} {
|
||||
t.Helper()
|
||||
if _, err := task.Store.AddIntent(map[string]any{"summary": "fetch the probe"}, 1, nil, "planner"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
intent, token := g.s.engine.claimNext(task, "work#1")
|
||||
if intent == nil {
|
||||
t.Fatal("no intent claimed")
|
||||
}
|
||||
worker := agent.NewWorker(p, "test-model", g.m.dir, nil, 0, 8, extra...)
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
g.s.engine.runIntent(context.Background(), task, "work#1", worker, intent, token, "", "")
|
||||
}()
|
||||
t.Cleanup(func() {
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(60 * time.Second):
|
||||
t.Error("the worker run did not end")
|
||||
}
|
||||
})
|
||||
return done
|
||||
}
|
||||
|
||||
// A gated task end to end: a worker's WebFetch waits on a request, a click
|
||||
// through the API pins the address and the call is made, the next call passes
|
||||
// on the pin, the pin is listed and revoked, and the call after asks again.
|
||||
func TestTargetGateEndToEnd(t *testing.T) {
|
||||
g := newGateServer(t)
|
||||
on := true
|
||||
task := g.newTask(t, db.TaskCreateOptions{TargetGate: &on})
|
||||
pg := g.m.pg
|
||||
srv, hits, host, port := countingServer(t)
|
||||
key := net.JoinHostPort(host, port)
|
||||
target := srv.URL + "/e2e"
|
||||
p := &toolScript{calls: fetches(target, target, target), results: make(chan string, 16),
|
||||
hold: map[int64]chan struct{}{3: make(chan struct{})}}
|
||||
done := runScriptedWorker(t, g, task, p)
|
||||
|
||||
var id int64
|
||||
waitUntil(t, "the first request", func() bool {
|
||||
var ok bool
|
||||
id, _, _, ok = openTargetRequest(t, pg, task.ID, key)
|
||||
return ok
|
||||
})
|
||||
if hits.Load() != 0 {
|
||||
t.Fatal("fetched before the approval")
|
||||
}
|
||||
if code, body := g.do(t, "POST", fmt.Sprintf("/api/intercept/pending/%d/decide", id), `{"decision":"allowed"}`); code != 200 {
|
||||
t.Fatalf("decide: %d %v", code, body)
|
||||
}
|
||||
if r := p.next(t); !strings.Contains(r, "probe-reached") || hits.Load() != 1 {
|
||||
t.Fatalf("first call after the approval: %q hits=%d", r, hits.Load())
|
||||
}
|
||||
if r := p.next(t); !strings.Contains(r, "probe-reached") || hits.Load() != 2 {
|
||||
t.Fatalf("second call, pinned: %q hits=%d", r, hits.Load())
|
||||
}
|
||||
var requests int
|
||||
if err := pg.QueryRow(`SELECT count(*) FROM intercept_pending WHERE task_id=$1 AND target_key IS NOT NULL`, task.ID).Scan(&requests); err != nil || requests != 1 {
|
||||
t.Fatalf("%d requests, %v; want the one that was approved", requests, err)
|
||||
}
|
||||
|
||||
code, body := g.do(t, "GET", "/api/tasks/"+task.ID+"/target-pins", "")
|
||||
pins, _ := body["pins"].([]any)
|
||||
if code != 200 || len(pins) != 1 {
|
||||
t.Fatalf("pins: %d %v", code, body)
|
||||
}
|
||||
pin := pins[0].(map[string]any)
|
||||
if pin["host"] != host || fmt.Sprint(pin["port"]) != port || pin["status"] != "active" || fmt.Sprint(pin["pinned_ips"]) != "["+host+"]" {
|
||||
t.Fatalf("listed pin = %v", pin)
|
||||
}
|
||||
pinID, _ := pin["id"].(float64)
|
||||
if code, _ := g.do(t, "POST", fmt.Sprintf("/api/tasks/%s/target-pins/%d/revoke", task.ID, int64(pinID)), ""); code != 200 {
|
||||
t.Fatalf("revoke: %d", code)
|
||||
}
|
||||
close(p.hold[3])
|
||||
|
||||
var again int64
|
||||
waitUntil(t, "the request after the revoke", func() bool {
|
||||
var ok bool
|
||||
again, _, _, ok = openTargetRequest(t, pg, task.ID, key)
|
||||
return ok && again != id
|
||||
})
|
||||
if code, body := g.do(t, "POST", fmt.Sprintf("/api/intercept/pending/%d/decide", again), `{"decision":"denied"}`); code != 200 {
|
||||
t.Fatalf("deny: %d %v", code, body)
|
||||
}
|
||||
if r := p.next(t); !strings.Contains(r, "TARGET_DENIED") || hits.Load() != 2 {
|
||||
t.Fatalf("third call, revoked then denied: %q hits=%d", r, hits.Load())
|
||||
}
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(60 * time.Second):
|
||||
t.Fatal("the worker run did not end")
|
||||
}
|
||||
}
|
||||
|
||||
// A custom tool in a gated worker session. A deny rule on the call decides
|
||||
// first, at the hook: the tool never runs, so its rendered target is never
|
||||
// read, resolved or requested. The next call, which no rule refuses, is gated
|
||||
// inside the tool on the URL it renders.
|
||||
func TestDenyRuleWinsBeforeRenderedTarget(t *testing.T) {
|
||||
g := newGateServer(t)
|
||||
on := true
|
||||
task := g.newTask(t, db.TaskCreateOptions{TargetGate: &on})
|
||||
pg, ic := g.m.pg, g.m.interceptor
|
||||
_, hits, _, port := countingServer(t)
|
||||
res := &countingResolver{answers: map[string][]string{"api.example.com": {"127.0.0.1"}}}
|
||||
ic.SetResolver(res)
|
||||
|
||||
prior, err := ic.GetEnabledTools()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = ic.SetEnabledTools(prior) })
|
||||
if err := ic.SetEnabledTools(append(slices.Clone(prior), "probe_http")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
const sentinel = "ARTEX_RENDERED_DENY_PROBE"
|
||||
rule, err := pg.CreateInterceptRule("rendered deny probe", "tool_input", "string", sentinel, "deny", "", 1000000, true, false, 0, "deny")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_ = pg.DeleteInterceptRule(rule.ID)
|
||||
ic.Invalidate()
|
||||
})
|
||||
ic.Invalidate()
|
||||
|
||||
tool := g.s.buildCustomTool(&db.Tool{Key: "probe_http", Kind: "http",
|
||||
Schema: json.RawMessage(`{"type":"object","properties":{"sub":{"type":"string"},"note":{"type":"string"}}}`),
|
||||
Exec: execJSON(t, map[string]any{"url": "http://{sub}.example.com:" + port + "/deny"})})
|
||||
p := &toolScript{calls: []scriptedCall{
|
||||
{"probe_http", `{"sub":"api","note":"` + sentinel + `"}`},
|
||||
{"probe_http", `{"sub":"api"}`},
|
||||
}, results: make(chan string, 16), hold: map[int64]chan struct{}{2: make(chan struct{})}}
|
||||
done := runScriptedWorker(t, g, task, p, tool)
|
||||
|
||||
if r := p.next(t); !strings.Contains(r, "intercepted by the platform") || strings.Contains(r, "TARGET_") {
|
||||
t.Fatalf("call refused by the rule: %q", r)
|
||||
}
|
||||
if n, rows := res.lookups(), targetRowCount(t, pg, task.ID); n != 0 || rows != 0 || hits.Load() != 0 {
|
||||
t.Fatalf("after the deny: %d lookups, %d target rows, %d requests sent; want none", n, rows, hits.Load())
|
||||
}
|
||||
close(p.hold[2])
|
||||
|
||||
key := "api.example.com:" + port
|
||||
var id int64
|
||||
waitUntil(t, "the request for the rendered target", func() bool {
|
||||
var ok bool
|
||||
var tool string
|
||||
id, tool, _, ok = openTargetRequest(t, pg, task.ID, key)
|
||||
return ok && tool == "probe_http"
|
||||
})
|
||||
if code, body := g.do(t, "POST", fmt.Sprintf("/api/intercept/pending/%d/decide", id), `{"decision":"denied"}`); code != 200 {
|
||||
t.Fatalf("deny: %d %v", code, body)
|
||||
}
|
||||
if r := p.next(t); !strings.Contains(r, "TARGET_DENIED") || hits.Load() != 0 || res.lookups() == 0 {
|
||||
t.Fatalf("rendered target: %q, %d requests sent, %d lookups", r, hits.Load(), res.lookups())
|
||||
}
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(60 * time.Second):
|
||||
t.Fatal("the worker run did not end")
|
||||
}
|
||||
}
|
||||
|
||||
// A task with the gate off (the default) runs as before: a worker's WebFetch
|
||||
// and the custom tools reach their targets, and the gate writes no row and
|
||||
// resolves no name.
|
||||
func TestGateOffTaskIsUnchanged(t *testing.T) {
|
||||
g := newGateServer(t)
|
||||
task := g.newTask(t, db.TaskCreateOptions{})
|
||||
res := &countingResolver{answers: map[string][]string{"localhost": {"127.0.0.1"}}}
|
||||
g.m.interceptor.SetResolver(res)
|
||||
_, hits, _, port := countingServer(t)
|
||||
byName := "http://localhost:" + port
|
||||
|
||||
p := &toolScript{calls: fetches(byName+"/a", byName+"/b"), results: make(chan string, 16)}
|
||||
done := runScriptedWorker(t, g, task, p)
|
||||
for n := 1; n <= 2; n++ {
|
||||
if r := p.next(t); !strings.Contains(r, "probe-reached") {
|
||||
t.Fatalf("WebFetch %d: %q", n, r)
|
||||
}
|
||||
}
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(60 * time.Second):
|
||||
t.Fatal("the worker run did not end")
|
||||
}
|
||||
|
||||
httpTool := g.s.buildCustomTool(&db.Tool{Key: "probe_http", Kind: "http",
|
||||
Schema: json.RawMessage(`{"type":"object","properties":{"port":{"type":"string"}}}`),
|
||||
Exec: execJSON(t, map[string]any{"url": "http://localhost:{port}/c"})})
|
||||
if out, isErr := callTool(t, httpTool, gatedCtx(task, intercept.DefaultTargetWaitBound), fmt.Sprintf(`{"port":%q}`, port), nil); isErr || !strings.Contains(out, "probe-reached") {
|
||||
t.Fatalf("custom http tool: error=%v out=%q", isErr, out)
|
||||
}
|
||||
cmdTool := g.s.buildCustomTool(&db.Tool{Key: "probe_cmd", Kind: "command",
|
||||
Exec: execJSON(t, map[string]any{"command": ": http://localhost:" + port + "/ && echo ran-gate-off"})})
|
||||
if out, isErr := callTool(t, cmdTool, gatedCtx(task, intercept.DefaultTargetWaitBound), `{}`, &actool.ToolContext{WorkingDir: t.TempDir()}); isErr || !strings.Contains(out, "ran-gate-off") {
|
||||
t.Fatalf("custom command tool: error=%v out=%q", isErr, out)
|
||||
}
|
||||
|
||||
if hits.Load() != 3 {
|
||||
t.Fatalf("%d requests reached the server, want 3", hits.Load())
|
||||
}
|
||||
if n := targetRowCount(t, g.m.pg, task.ID); n != 0 {
|
||||
t.Fatalf("%d target rows, want none", n)
|
||||
}
|
||||
if n := res.lookups(); n != 0 {
|
||||
t.Fatalf("%d lookups, want none", n)
|
||||
}
|
||||
}
|
||||
@@ -128,7 +128,7 @@ func TestRunHTTPTool(t *testing.T) {
|
||||
"headers": map[string]string{"Authorization": "Bearer {token}"},
|
||||
"body": `{"flag":"{flag}"}`,
|
||||
})
|
||||
res, err := (&Server{}).runHTTPTool(context.Background(), execRaw,
|
||||
res, err := (&Server{}).runHTTPTool(context.Background(), "probe", execRaw,
|
||||
map[string]any{"flag": "CTF{x}", "token": "sekret"}, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("runHTTPTool: %v", err)
|
||||
@@ -176,7 +176,7 @@ func TestHTTPToolTruncatesLargeBody(t *testing.T) {
|
||||
defer srv.Close()
|
||||
|
||||
execRaw, _ := json.Marshal(map[string]any{"method": "GET", "url": srv.URL})
|
||||
res, err := (&Server{}).runHTTPTool(context.Background(), execRaw, map[string]any{}, nil)
|
||||
res, err := (&Server{}).runHTTPTool(context.Background(), "probe", execRaw, map[string]any{}, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("runHTTPTool: %v", err)
|
||||
}
|
||||
|
||||
+23
-1
@@ -13,6 +13,7 @@ import (
|
||||
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/intercept"
|
||||
"github.com/Autumn-27/norma/permission"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
)
|
||||
@@ -253,7 +254,7 @@ func (s *Server) toolSpawnTask() actool.CoreTool {
|
||||
"plan_heartbeat_seconds": map[string]any{"type": "integer", "description": "optional: planner heartbeat interval (seconds). When this much time has passed since the last planning round ended/the task started with no trigger in between → run one planning round (deadlock backstop + wakes the planner to supervise in-flight workers). Empty or 0 = default 600 (10min);"},
|
||||
"seed_first_intent": map[string]any{"type": "boolean", "description": "optional: enable for simple tasks to issue one seed intent at creation (content = description + goal) so a worker can start testing right away without waiting for the first planner round; default false (the standard plan-then-execute flow)."},
|
||||
}, "description", "goal"),
|
||||
func(_ context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
func(ctx context.Context, in json.RawMessage) (actool.Result, error) {
|
||||
var a struct {
|
||||
Description string `json:"description"`
|
||||
Goal string `json:"goal"`
|
||||
@@ -307,11 +308,32 @@ func (s *Server) toolSpawnTask() actool.CoreTool {
|
||||
if pin != nil {
|
||||
llmIDs = []int64{*pin}
|
||||
}
|
||||
// Target gate: a child spawned from a gated scope is gated too, and keeps
|
||||
// the block rules of the calling task. Read from the caller's scope, never
|
||||
// from parent_ref (optional and written by the model). Pins and allow rules
|
||||
// are never copied. Gate-off caller or no scope: created as before.
|
||||
var gate *bool
|
||||
var blockRules []db.TaskInterceptRuleInput
|
||||
if scope, ok := intercept.TargetScopeFromContext(ctx); ok && scope.Enabled != nil && scope.Enabled() {
|
||||
on := true
|
||||
gate = &on
|
||||
if callerID, err := strconv.ParseInt(scope.TaskID, 10, 64); err == nil && callerID > 0 {
|
||||
block, _, err := s.m.Assets().TaskInterceptRulesSplit(callerID)
|
||||
if err != nil {
|
||||
return actool.Errorf("read the block rules of the calling task: " + err.Error()), nil
|
||||
}
|
||||
for _, r := range block {
|
||||
blockRules = append(blockRules, db.TaskInterceptRuleInput{Enabled: r.Enabled, Action: r.Action, Kind: r.Kind, Pattern: r.Pattern, Note: r.Note})
|
||||
}
|
||||
}
|
||||
}
|
||||
t, err := s.m.CreateTaskWithOptions(a.Description, a.Goal, db.TaskCreateOptions{
|
||||
SourceTaskIDs: sourceIDs,
|
||||
LLMProfileIDs: llmIDs,
|
||||
TimeoutSeconds: a.TimeoutSeconds,
|
||||
PlanHeartbeatSeconds: a.PlanHeartbeatSeconds,
|
||||
InterceptRules: blockRules,
|
||||
TargetGate: gate,
|
||||
})
|
||||
if err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
|
||||
@@ -147,6 +147,9 @@ type triggeredRun struct {
|
||||
taskDesc string // task description (task-level, same across one task); rendered once when merged
|
||||
taskGoal string // task goal (task-level, same across one task); rendered once when merged
|
||||
mergeable bool // true for finding/goal event triggers (merge by taskID)
|
||||
// targetGate: the source task had the target gate on when the event fired; a
|
||||
// merged run is gated if any of its fires is. Never carries pins.
|
||||
targetGate bool
|
||||
}
|
||||
|
||||
func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDir string, opts ...Option) *Server {
|
||||
|
||||
+246
-21
@@ -5,16 +5,19 @@ import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"iter"
|
||||
"net/http/httptest"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/intercept"
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
)
|
||||
|
||||
// Target-gate opt-in flag, pins and decisions over the real mux and the real
|
||||
@@ -180,39 +183,261 @@ func TestCreateTaskAcceptsTargetGate(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// 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)
|
||||
// ruleLines renders a task's intercept rules for comparison.
|
||||
func ruleLines(t *testing.T, as *db.AssetStore, taskID string) []string {
|
||||
t.Helper()
|
||||
n, _ := strconv.ParseInt(taskID, 10, 64)
|
||||
rules, err := as.ListTaskInterceptRules(n)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
childID, found := strings.CutPrefix(res.Flatten(), "task created: ")
|
||||
if !found {
|
||||
t.Fatalf("spawn_task result: %q", res.Flatten())
|
||||
var out []string
|
||||
for _, r := range rules {
|
||||
out = append(out, fmt.Sprintf("%s %s %s enabled=%v %q", r.Action, r.Kind, r.Pattern, r.Enabled, r.Note))
|
||||
}
|
||||
g.cleanupTask(t, childID)
|
||||
return out
|
||||
}
|
||||
|
||||
// checkSpawnedChild checks the target-gate state of a spawned task: the flag,
|
||||
// live and stored, no pin, and its own intercept rules.
|
||||
func (g *gateServer) checkSpawnedChild(t *testing.T, childID string, gate bool, rules []string) {
|
||||
t.Helper()
|
||||
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 !ok {
|
||||
t.Fatalf("spawned task %s not found", childID)
|
||||
}
|
||||
if child.targetScope(0).Enabled() {
|
||||
t.Fatal("the spawned task inherited the gate")
|
||||
if child.targetScope(0).Enabled() != gate {
|
||||
t.Fatalf("spawned task gate = %v, want %v", !gate, 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 stored, err := g.m.pg.GetTask(n); err != nil || stored == nil || stored.TargetGate != gate {
|
||||
t.Fatalf("stored child = %+v, %v; want the gate %v", stored, err, gate)
|
||||
}
|
||||
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)
|
||||
}
|
||||
if got := ruleLines(t, g.m.Assets(), childID); !slices.Equal(got, rules) {
|
||||
t.Fatalf("spawned task rules = %q, want %q", got, rules)
|
||||
}
|
||||
}
|
||||
|
||||
// A task spawned through spawn_task takes the target gate of the calling scope
|
||||
// (D23) and, from a gated task, that task's block rules (D24). It never takes
|
||||
// pins or allow rules, and parent_ref plays no part. A gate-off caller and a
|
||||
// call with no scope create the child as before: gate off, no rules.
|
||||
func TestSpawnedTaskInheritsGateFromCaller(t *testing.T) {
|
||||
g := newGateServer(t)
|
||||
on := true
|
||||
parentRules := []db.TaskInterceptRuleInput{
|
||||
{Enabled: true, Action: "block", Kind: "exact_domain", Pattern: "blocked.example.com", Note: "out of scope"},
|
||||
{Enabled: false, Action: "block", Kind: "cidr", Pattern: "198.51.100.0/24"},
|
||||
{Enabled: true, Action: "allow", Kind: "exact_domain", Pattern: "allowed.example.com"},
|
||||
}
|
||||
gated := g.newTask(t, db.TaskCreateOptions{TargetGate: &on, InterceptRules: parentRules})
|
||||
open := g.newTask(t, db.TaskCreateOptions{InterceptRules: parentRules})
|
||||
if !gated.targetScope(0).Enabled() || open.targetScope(0).Enabled() {
|
||||
t.Fatal("control: one parent gated, the other not")
|
||||
}
|
||||
// The gated parent holds one pin, which no child may get.
|
||||
if p, ok, err := g.m.pg.ResolveTargetRequest(openTargetRow(t, g.m.pg, taskRef(t, gated), "app.example.com", "203.0.113.7"), true, "ARTEX", "192.0.2.1"); err != nil || !ok || p == nil {
|
||||
t.Fatalf("pin the parent: %+v ok=%v err=%v", p, ok, err)
|
||||
}
|
||||
conv, err := g.m.pg.CreateConversationGated("auto", "spawn from a gated chat", nil, true)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = g.m.pg.DeleteConversation(conv.ID) })
|
||||
gatedChat := intercept.TargetScope{ConversationID: conv.ID, Enabled: func() bool { return true },
|
||||
WaitBound: intercept.DefaultTargetWaitBound, Interceptor: g.m.interceptor}
|
||||
scoped := func(s intercept.TargetScope) context.Context {
|
||||
return intercept.WithTargetScope(context.Background(), s)
|
||||
}
|
||||
blockOnly := []string{
|
||||
`block exact_domain blocked.example.com enabled=true "out of scope"`,
|
||||
`block cidr 198.51.100.0/24 enabled=false ""`,
|
||||
}
|
||||
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
ctx context.Context
|
||||
parentRef string
|
||||
gate bool
|
||||
rules []string
|
||||
}{
|
||||
{"gated task", scoped(gated.targetScope(intercept.DefaultTargetWaitBound)), "", true, blockOnly},
|
||||
{"gated task, parent_ref names a gate-off task", scoped(gated.targetScope(0)), open.ID, true, blockOnly},
|
||||
{"gated conversation", scoped(gatedChat), "", true, nil},
|
||||
{"gate-off task", scoped(open.targetScope(0)), "", false, nil},
|
||||
{"gate-off task, parent_ref names a gated task", scoped(open.targetScope(0)), gated.ID, false, nil},
|
||||
{"no scope, parent_ref names a gated task", context.Background(), gated.ID, false, nil},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
args := map[string]any{"description": "spawned child", "goal": gateTestGoal}
|
||||
if tc.parentRef != "" {
|
||||
args["parent_ref"] = tc.parentRef
|
||||
}
|
||||
in, _ := json.Marshal(args)
|
||||
res, err := g.s.toolSpawnTask().Call(tc.ctx, 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)
|
||||
if child, _ := g.m.Task(childID); child == nil || child.ParentRef != tc.parentRef {
|
||||
t.Fatalf("spawned task = %+v, want parent_ref %q", child, tc.parentRef)
|
||||
}
|
||||
g.checkSpawnedChild(t, childID, tc.gate, tc.rules)
|
||||
})
|
||||
}
|
||||
// The callers are unchanged.
|
||||
if pins, err := g.m.pg.ListTargetPins(taskRef(t, gated)); err != nil || len(pins) != 1 {
|
||||
t.Fatalf("parent pins = %+v, %v; want still one", pins, err)
|
||||
}
|
||||
if got := ruleLines(t, g.m.Assets(), gated.ID); len(got) != 3 {
|
||||
t.Fatalf("parent rules = %q, want the three it was created with", got)
|
||||
}
|
||||
}
|
||||
|
||||
// toolCallProvider is a model that makes one tool call on its first turn and
|
||||
// ends the run on the next.
|
||||
type toolCallProvider struct {
|
||||
tool, input string
|
||||
turn atomic.Int64
|
||||
}
|
||||
|
||||
func (p *toolCallProvider) Stream(context.Context, llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
|
||||
events := []llm.StreamEvent{{Type: llm.SETextDelta, Text: "done"}, {Type: llm.SEMessageDelta, StopReason: "end_turn"}}
|
||||
if p.turn.Add(1) == 1 {
|
||||
events = []llm.StreamEvent{
|
||||
{Type: llm.SEToolUseStart, ToolID: "call-1", ToolName: p.tool},
|
||||
{Type: llm.SEToolInputJSON, Text: p.input},
|
||||
{Type: llm.SEMessageDelta, StopReason: "tool_use"},
|
||||
}
|
||||
}
|
||||
return func(yield func(llm.StreamEvent, error) bool) {
|
||||
for _, event := range events {
|
||||
if !yield(event, nil) {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (p *toolCallProvider) Complete(ctx context.Context, req llm.CompletionRequest) (llm.Message, string, llm.Usage, error) {
|
||||
return accumulateStreamForTest(ctx, p.Stream, req)
|
||||
}
|
||||
|
||||
// spawn_task called by the model of an auto conversation, through the real
|
||||
// harness, tool binding and tool context: the child takes the conversation's
|
||||
// flag. The unit test above builds the caller's context by hand; this one shows
|
||||
// the conversation's scope reaches the tool.
|
||||
func TestSpawnTaskSeesCallerScopeThroughHarness(t *testing.T) {
|
||||
g := newGateServer(t)
|
||||
for _, gate := range []bool{true, false} {
|
||||
t.Run(fmt.Sprintf("conversation gate %v", gate), func(t *testing.T) {
|
||||
desc := fmt.Sprintf("harness child %v %d", gate, time.Now().UnixNano())
|
||||
in, _ := json.Marshal(map[string]any{"description": desc, "goal": gateTestGoal})
|
||||
conv, err := g.m.pg.CreateConversationGated("auto", "spawn through the harness", nil, gate)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = g.m.pg.DeleteConversation(conv.ID) })
|
||||
g.s.cfgMu.Lock()
|
||||
g.s.chatAgent = agent.NewChatAgent(&toolCallProvider{tool: "spawn_task", input: string(in)}, "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(), conv.ID))
|
||||
g.s.runConversationTurn(ctx, cancel, conv, "spawn the child task", g.s.convBusyKey(conv.ID))
|
||||
var childID string
|
||||
for _, task := range g.m.List() {
|
||||
if task.Description == desc {
|
||||
childID = task.ID
|
||||
}
|
||||
}
|
||||
if childID == "" {
|
||||
t.Fatal("the model's spawn_task call created no task")
|
||||
}
|
||||
g.cleanupTask(t, childID)
|
||||
g.checkSpawnedChild(t, childID, gate, nil)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// A conversation started by a trigger of a gated task is gated (D25): a single
|
||||
// fire, a by-task merge, and a merge across tasks of which one is gated. A fire
|
||||
// of a gate-off task, or with no task, starts it gate off as before. The flag is
|
||||
// copied, never pins.
|
||||
func TestTriggeredConversationInheritsGate(t *testing.T) {
|
||||
g := newGateServer(t)
|
||||
g.s.cfgMu.Lock()
|
||||
g.s.chatAgent = agent.NewChatAgent(newScopeProbeProvider(), "test-model", g.m.dir, nil, 10000)
|
||||
g.s.cfgMu.Unlock()
|
||||
on := true
|
||||
gated, open := g.newTask(t, db.TaskCreateOptions{TargetGate: &on}), g.newTask(t, db.TaskCreateOptions{})
|
||||
if p, ok, err := g.m.pg.ResolveTargetRequest(openTargetRow(t, g.m.pg, taskRef(t, gated), "app.example.com", "203.0.113.7"), true, "ARTEX", "192.0.2.1"); err != nil || !ok || p == nil {
|
||||
t.Fatalf("pin the gated task: %+v ok=%v err=%v", p, ok, err)
|
||||
}
|
||||
gatedID, openID := taskRef(t, gated).TaskID, taskRef(t, open).TaskID
|
||||
|
||||
// The server context is cancelled, so StartTriggeredRun only queues. The test
|
||||
// takes the next run with the merge mode under test and runs it as the pump would.
|
||||
trigger := func(t *testing.T, mode string, taskIDs ...int64) *db.Conversation {
|
||||
t.Helper()
|
||||
key := fmt.Sprintf("gate-trigger-%d", time.Now().UnixNano())
|
||||
for _, id := range taskIDs {
|
||||
g.s.StartTriggeredRun(key, "trigger", "event body", id, id > 0, "desc", "goal")
|
||||
}
|
||||
g.s.queueMu.Lock()
|
||||
item := g.s.nextTriggerRun(key, triggerBehavior{runMode: "serial", mergeMode: mode})
|
||||
left := len(g.s.triggerQ[key])
|
||||
delete(g.s.triggerQ, key)
|
||||
delete(g.s.triggerCfg, key)
|
||||
g.s.queueMu.Unlock()
|
||||
if left != 0 {
|
||||
t.Fatalf("%d fires left in the queue, want one run", left)
|
||||
}
|
||||
item.title = key
|
||||
g.s.runTriggeredRun(item)
|
||||
convs, err := g.m.pg.ListConversations()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, c := range convs {
|
||||
if c.Title == key {
|
||||
t.Cleanup(func() { _ = g.m.pg.DeleteConversation(c.ID) })
|
||||
return c
|
||||
}
|
||||
}
|
||||
t.Fatal("the trigger created no conversation")
|
||||
return nil
|
||||
}
|
||||
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
mode string
|
||||
tasks []int64
|
||||
gate bool
|
||||
}{
|
||||
{"gated task", "none", []int64{gatedID}, true},
|
||||
{"gated task, merged by task", "by_task", []int64{gatedID, gatedID}, true},
|
||||
{"merge across tasks, gated second", "all", []int64{openID, gatedID}, true},
|
||||
{"gate-off task", "none", []int64{openID}, false},
|
||||
{"gate-off task, merged by task", "by_task", []int64{openID, openID}, false},
|
||||
{"no task", "none", []int64{0}, false},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
c := trigger(t, tc.mode, tc.tasks...)
|
||||
if c.TargetGate != tc.gate {
|
||||
t.Fatalf("triggered conversation gate = %v, want %v", c.TargetGate, tc.gate)
|
||||
}
|
||||
if pins, err := g.m.pg.ListTargetPins(db.TargetScopeRef{ConversationID: c.ID}); err != nil || len(pins) != 0 {
|
||||
t.Fatalf("the triggered conversation has pins: %+v, %v", pins, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// The conversation flag is set on create and on patch, and the turn that starts
|
||||
|
||||
Reference in New Issue
Block a user