mirror of
https://github.com/BaaRai/ARTEX.git
synced 2026-11-16 05:00:53 +08:00
Compare commits
21
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ce096cdba1 | ||
|
|
50f6d07640 | ||
|
|
e6ab0f34a7 | ||
|
|
5ed28cde9f | ||
|
|
773fa9183a | ||
|
|
1cbab7cc39 | ||
|
|
bc75b5dc43 | ||
|
|
cddcd98f3b | ||
|
|
d42988b6a8 | ||
|
|
3975e58ce0 | ||
|
|
fdba29f1a9 | ||
|
|
a59d430175 | ||
|
|
c2d1f557c3 | ||
|
|
e636408f66 | ||
|
|
fcbd87445e | ||
|
|
59f768f1ef | ||
|
|
7f3edf9bfe | ||
|
|
2e55615713 | ||
|
|
251ab680cc | ||
|
|
ac6e2b3c44 | ||
|
|
97b0fdced6 |
@@ -29,6 +29,10 @@ jobs:
|
||||
- uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version-file: go.mod
|
||||
- name: gofmt
|
||||
run: |
|
||||
files=$(gofmt -l $(git ls-files '*.go'))
|
||||
if [ -n "$files" ]; then echo "gofmt required:"; echo "$files"; exit 1; fi
|
||||
- run: go build ./...
|
||||
- run: go vet ./...
|
||||
# Skipped tests, tracked in issue #9:
|
||||
|
||||
+34
-1
@@ -4,9 +4,42 @@ All notable changes to this project are documented in this file. The format is b
|
||||
|
||||
## [Unreleased]
|
||||
|
||||
### Added
|
||||
|
||||
- **Target noise governor**: one shared egress governor paces ARTEX's own HTTP requests (WebFetch, the http custom tools and the enrich probe) against a noise profile (`off` by default, `quiet`, `balanced`, `aggressive`): a per-host token bucket with jitter, per-host and global in-flight caps, 429/503 backoff, GET deduplication with a short cache, and a configurable target User-Agent (`ARTEX_TARGET_USER_AGENT`). Off is a strict passthrough (#205).
|
||||
- **Low-and-slow prompt block**: when a noise profile is set, the planner, worker, main agent and chat prompts carry the per-host limits of the profile and the exact rate or delay flags for nmap, nuclei, httpx, katana, ffuf, dalfox and sqlmap; with the profile off the prompts are unchanged (#212).
|
||||
- **Noise profile at runtime**: System settings has a **Target noise** card. It sets the profile (`off`, `quiet`, `balanced`, `aggressive`) through the `noise_profile` setting, which takes precedence over `ARTEX_NOISE_PROFILE`. The change applies to the next WebFetch, http custom tool and enrich request, and to the next agent prompt, without a restart. The card also shows the per-host counters, which `GET /api/settings/noise` returns (#233).
|
||||
- **Bash requests paced through the capture proxy**: with traffic capture on, the Bash proxy URL carries a marker userinfo. The recording proxy paces each request of a marked connection by the noise profile's per-host rate and jitter, and counts it, before it goes upstream. This applies inside CONNECT tunnels too. WebFetch, the http custom tools and enrich keep the unmarked URL and are not paced twice. The proxy now drops `Proxy-Authorization` from every request it forwards or records. This is best effort: only clients that honour the proxy are covered, with capture on (#234).
|
||||
- **Hierarchical budgets**, off by default: a task budget and a per-intent budget take the minimum of their own budget and the parent's remaining budget above the wrap-up reserve, planner calls count inside the task, and reaching the reserve forces a wrap-up (#206).
|
||||
- **Coverage ledger**: the worker records what it examined, ruled out or left open for each surface and risk class with `record_coverage`; gaps, unexamined classes and completeness appear in the planner context, `GET /api/tasks/{id}/coverage`, the coverage card and the Markdown report (#207).
|
||||
- **Closed-question judge**: the LLM fallback judge answers eight fixed questions that the code owns and the code takes the decision; every failure now leads to ask instead of allow, and the previous behaviour stays available as `llm_judge_mode = freeform` (#208).
|
||||
- **argv interception rules**: a new `match_type = argv` matches the parsed prefix of a Bash command exactly, and a pending Bash approval can create such a rule in one click; a code-owned denylist refuses programs where a prefix would allow arbitrary code (#210, #211).
|
||||
- **Scanner output ingestion**: `ingest_scan_output` parses nmap XML, nuclei JSONL, httpx JSONL and ffuf JSON, and writes assets and facts deterministically, so the model never reads the raw output (#214, #215).
|
||||
- **bbot ingestion**: a typed parser splits bbot JSONL events by `scope_distance`, only in-scope events go through `insert_assets`, and `ingest_bbot` reads the file line by line with a 64 KiB cap (#199, #200).
|
||||
- **Tool RAG**, off by default: when the bound tool list of an agent is large, only the subset relevant to the current intent is injected, as a projection that never adds or edits a tool (`ARTEX_TOOL_RAG`) (#204).
|
||||
- **Fact provenance**: `record_fact` takes a closed `source` enum (default `inference`) and every fact carries the server-set `origin`, both shown in `recent_facts`, `list_facts` and `expand_digest` (#194).
|
||||
- **Per-message nonce envelope**: agent- and target-derived fields reach the planner and worker user messages inside a code-owned data envelope with a random nonce per message; operator input stays outside it (#195).
|
||||
- **Intercept risk score and batch approval**: a deterministic risk score (low, medium, high) is shown on each pending approval, and up to 100 pending decisions can be settled in one call, each through the existing single-decision path (#232).
|
||||
- **`ask_operator` target warning**: a request whose question names a host or URL outside the goal targets is tagged as containing target-derived text (#197).
|
||||
- **Opt-in update banner**: `ARTEX_UPDATE_CHECK` makes the backend check the latest GitHub release at most once per 24 hours and the UI show a banner; off by default (#202).
|
||||
- **Release provenance**: release archives carry GitHub build provenance attestations and a CycloneDX SBOM (`gh attestation verify <zip> --repo BaaRai/ARTEX`); release notes are the matching `CHANGELOG.md` section (#219, #201).
|
||||
|
||||
### Changed
|
||||
|
||||
- **Untrusted-output boundary is now on by default**: Norma frames tool results for the worker, planner, chat and main agent unless `ARTEX_UNTRUSTED_BOUNDARY` is set to `0`, `false`, `off` or `no`. It was opt-in via `ARTEX_UNTRUSTED_BOUNDARY=1` in 0.6.0 (#166).
|
||||
- **Untrusted-output boundary is now on by default**: Norma frames tool results for the worker, planner, chat and main agent unless `ARTEX_UNTRUSTED_BOUNDARY` is set to `0`, `false`, `off` or `no`. It was opt-in via `ARTEX_UNTRUSTED_BOUNDARY=1` in 0.6.0 (#166); the default flips here (#203).
|
||||
- **Versioned migrations**: the `db` package applies numbered `db/migrations/NNNN_slug.sql` files once, in order, under the schema advisory lock, and records a SHA-256 checksum in `schema_migrations`; missing, renamed, changed or out-of-order migrations fail closed (#226).
|
||||
- **Coverage graph performance**: the coverage graph batches its connector lookups (three queries instead of one per root and per company) and adds the partial index `idx_ts_icp`; on a task of about 3,000 assets the p50 dropped from 126 ms to 27 ms and the queries per call from 221 to 3 (#230).
|
||||
|
||||
### Fixed
|
||||
|
||||
- **Auto-scope hygiene**: auto `task_scope` rows count for coverage and untested assets only under a declared root, and direct coverage counts anchored assets through fact anchors only; the scope is described as "coverage scope, not enforced" (#196).
|
||||
- **Enrich HTTP probe**: it now dials through `intercept.PinnedDialer`, so the target gate applies and a probe without a target scope is refused (#198).
|
||||
- **Agent hook data races**: the eight package-level hooks are synchronised, so a second wire-up no longer races tasks that already read them (#227).
|
||||
|
||||
### Security
|
||||
|
||||
- **Secrets at rest**: the credential-bearing columns and settings (LLM API keys and proxy, task credentials, MCP server env, notification channel config, search and proxy settings), the exec spec of custom tools and the evidence snapshot heads are encrypted in PostgreSQL with AES-256-GCM (`enc:v1:`) using `ARTEX_SECRET_KEY`, else `keys/secret.key`; existing plaintext rows are backfilled once at startup, and a missing or wrong key while encrypted data exists refuses to start. This protects database dumps and backups, not a running instance (#213, #216, #217, #218).
|
||||
- **URL credentials are dropped before storage**: scan and bbot ingestion, the asset store and archive restore strip URL userinfo, so `https://user:pass@host/app` is stored and matched as `https://host/app` (#222, #223, #224).
|
||||
|
||||
## [0.6.0] - 2026-11-12
|
||||
|
||||
|
||||
@@ -6,6 +6,8 @@
|
||||
|
||||
**Autonomous penetration testing that keeps its receipts.**
|
||||
|
||||
(upstream: [Autumn-27/ARTEX](https://github.com/Autumn-27), AGPL-3.0)
|
||||
|
||||
Self-hosted multi-agent offensive security. A planner explores, workers run real tools,
|
||||
and every finding stays bound to the assets and HTTP traffic that prove it.
|
||||
|
||||
@@ -405,8 +407,15 @@ flight), `aggressive` = 20 req/s (burst 40, 4 in flight). A profile adds jitter,
|
||||
429/503 (honours `Retry-After`, bounded retries, GET/HEAD/OPTIONS only), deduplicates identical
|
||||
in-flight GETs and caches them briefly (60/30/10 s; never POST/PUT/PATCH/DELETE), and logs a
|
||||
per-host noise line each minute. `ARTEX_TARGET_USER_AGENT` sets a User-Agent on target requests
|
||||
that carry none. Bash-launched tools are not paced yet; the low-and-slow prompt block supplies
|
||||
their rate flags.
|
||||
that carry none. With traffic capture on, Bash commands that use the proxy (curl, wget,
|
||||
python-requests) are paced too, per host by the same rate and jitter, and counted in the same
|
||||
per-host lines. The in-flight caps, retries and cache apply only to ARTEX's own requests. This is
|
||||
best effort, not a security boundary. Raw scanners such as nmap are not paced, and neither is a
|
||||
command that bypasses the proxy (`--noproxy`, a removed proxy variable). With capture off,
|
||||
nothing launched from Bash is paced. The low-and-slow prompt block supplies the rate flags in
|
||||
every case.
|
||||
The **Target noise** card in System settings sets the profile at runtime (no restart) and shows the
|
||||
per-host counters; once saved, that setting takes precedence over `ARTEX_NOISE_PROFILE`.
|
||||
|
||||
**Low-and-slow agent guidance (off by default)**: when a noise profile is set, the planner, worker,
|
||||
main agent and chat agents receive a code-owned prompt block with the profile's per-host limits and
|
||||
@@ -686,6 +695,12 @@ Operational security properties:
|
||||
|
||||
---
|
||||
|
||||
## About this fork
|
||||
|
||||
This repository is the English translation and continuation of ARTEX, originally by [Autumn-27](https://github.com/Autumn-27). The upstream repository is no longer public, so the link points to the author profile. Changes in this fork are tracked in [CHANGELOG.md](CHANGELOG.md), under the same AGPL-3.0 license.
|
||||
|
||||
---
|
||||
|
||||
## References and prior art
|
||||
|
||||
- [Cairn](https://github.com/oritera/Cairn)
|
||||
|
||||
+8
-2
@@ -169,9 +169,12 @@ func (c *ChatAgent) Chat(ctx context.Context, agentKey, sessionID, message strin
|
||||
if c.guard != nil {
|
||||
opts.Hooks = c.guard.Hooks()
|
||||
}
|
||||
// The session id is the provider-side identity (session header, usage metering)
|
||||
// with or without persistence; Norma only puts it on ctx when it writes a transcript.
|
||||
opts.SessionID = sessionID
|
||||
ctx = transcript.WithSessionID(ctx, opts.SessionID)
|
||||
if c.tx != nil { // persist raw human↔AI conversation; one accumulating file per thread
|
||||
opts.Transcript = c.tx
|
||||
opts.SessionID = sessionID
|
||||
}
|
||||
// experimental feature: when enabled, noa takes over context compression (archive centralized under <workDir>/noa/<SessionID>, persistent).
|
||||
enableNoa(&opts, c.noaEnabledFn, c.workDir, "chat-"+sessionID, noaWarn("chat-"+sessionID))
|
||||
@@ -180,8 +183,11 @@ func (c *ChatAgent) Chat(ctx context.Context, agentKey, sessionID, message strin
|
||||
defer s.Close()
|
||||
// reload prior conversation so the agent has context across turns (each Chat is
|
||||
// a fresh session). First turn: no file yet → Resume loads nothing and proceeds.
|
||||
// A corrupt or unreadable transcript stops the turn before any model or tool call.
|
||||
if c.tx != nil {
|
||||
_ = s.Resume(sessionID)
|
||||
if err := resumeSession(s, opts.SessionID); err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
// re-unlock skill-gated MCPs from prior Skill() calls in the reloaded history so
|
||||
// revealed tools stay callable across the fresh session.
|
||||
|
||||
+13
-7
@@ -169,14 +169,17 @@ func (m *MainAgent) Chat(ctx context.Context, taskID int64, mainSeg int, as *db.
|
||||
if hooks != nil { // typed-nil guard: only set when concrete (avoids harness panic)
|
||||
opts.Hooks = hooks
|
||||
}
|
||||
// Segment 0 keeps the legacy "exp%d-main" name so existing transcripts still
|
||||
// load; each new session (seg>=1) gets its own file for a clean context. The id
|
||||
// is the provider-side identity (session header, usage metering) with or without
|
||||
// persistence; Norma only puts it on ctx when it writes a transcript.
|
||||
opts.SessionID = fmt.Sprintf("exp%d-main", ts.ID())
|
||||
if mainSeg > 0 {
|
||||
opts.SessionID = fmt.Sprintf("exp%d-main-s%d", ts.ID(), mainSeg)
|
||||
}
|
||||
ctx = transcript.WithSessionID(ctx, opts.SessionID)
|
||||
if m.tx != nil { // persist raw human↔AI conversation; one accumulating file per segment
|
||||
opts.Transcript = m.tx
|
||||
// Segment 0 keeps the legacy "exp%d-main" name so existing transcripts still
|
||||
// load; each new session (seg>=1) gets its own file for a clean context.
|
||||
opts.SessionID = fmt.Sprintf("exp%d-main", ts.ID())
|
||||
if mainSeg > 0 {
|
||||
opts.SessionID = fmt.Sprintf("exp%d-main-s%d", ts.ID(), mainSeg)
|
||||
}
|
||||
}
|
||||
// experimental feature: when enabled, noa takes over context compression (archive centralized under <workDir>/noa/<SessionID>, persistent).
|
||||
// the session id follows the same rule as the transcript (segment-aware), keeping archive and resume aligned.
|
||||
@@ -191,8 +194,11 @@ func (m *MainAgent) Chat(ctx context.Context, taskID int64, mainSeg int, as *db.
|
||||
// reload the prior conversation from the transcript so the agent has context
|
||||
// across turns (each Chat is a fresh session; without this it can't see earlier
|
||||
// messages). First turn: no file yet → Resume loads nothing and proceeds.
|
||||
// A corrupt or unreadable transcript stops the turn before any model or tool call.
|
||||
if m.tx != nil {
|
||||
_ = s.Resume(opts.SessionID)
|
||||
if err := resumeSession(s, opts.SessionID); err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
// C2: this session is fresh each turn; re-unlock skill-gated MCPs from prior
|
||||
// Skill() calls in the reloaded history so revealed tools stay callable.
|
||||
|
||||
@@ -14,10 +14,11 @@ import (
|
||||
// prompts stay byte-for-byte unchanged by default.
|
||||
const noiseProfileHeader = "[Noise profile: "
|
||||
|
||||
// noiseProfileBlock returns the low-and-slow guidance for the active noise
|
||||
// profile (ARTEX_NOISE_PROFILE), or "" when the profile is off. The text is
|
||||
// code-owned and carries no target data.
|
||||
func noiseProfileBlock() string { return renderNoiseProfile(egress.ProfileFromEnv()) }
|
||||
// noiseProfileBlock returns the low-and-slow guidance for the current noise
|
||||
// profile (the noise_profile setting, else ARTEX_NOISE_PROFILE), or "" when the
|
||||
// profile is off. It is read at each prompt build, so a changed setting shows in
|
||||
// the next prompt. The text is code-owned and carries no target data.
|
||||
func noiseProfileBlock() string { return renderNoiseProfile(egress.CurrentProfile()) }
|
||||
|
||||
// renderNoiseProfile renders the block for one profile. It is pure: tests cover
|
||||
// every profile without touching the environment.
|
||||
@@ -63,6 +64,6 @@ func renderNoiseProfile(p egress.Profile) string {
|
||||
fmt.Fprintf(&b, "- sqlmap: --delay=%s --threads=1\n", delaySec)
|
||||
fmt.Fprintf(&b, "For any other tool, stay at or below %d requests per second per host, at most %s in flight, and at least %d ms between requests, using the tool's own rate or delay flag. Keep bbot within the same limits with its own rate settings, and do not enable active modules the program policy does not allow.\n", rps, inFlight, delayMS)
|
||||
b.WriteString("Prefer passive and read-only checks; test one host at a time; do not run two scanners against the same host in parallel; on 429 or 503 stop and back off; never raise the rate to meet a deadline.\n")
|
||||
b.WriteString("ARTEX paces its own WebFetch, http custom tool and enrich traffic with this profile, but commands launched through Bash are not paced: these flags are the only brake.\n")
|
||||
b.WriteString("ARTEX paces its own WebFetch, http custom tool and enrich traffic with this profile. With traffic capture on, Bash commands that use the capture proxy (curl, wget, python-requests) are paced too, by rate and jitter; raw scanners such as nmap and commands that bypass the proxy are not. These flags remain the brake: do not rely on the proxy.\n")
|
||||
return b.String()
|
||||
}
|
||||
|
||||
+16
-11
@@ -41,7 +41,7 @@ func TestRenderNoiseProfileTable(t *testing.T) {
|
||||
fmt.Sprintf("- dalfox scan: --rate-limit %d --workers %d", c.rps, c.conc),
|
||||
fmt.Sprintf("- sqlmap: --delay=%s --threads=1", c.sql),
|
||||
fmt.Sprintf("at least %d ms between requests", c.delay),
|
||||
"commands launched through Bash are not paced",
|
||||
"ARTEX paces its own WebFetch, http custom tool and enrich traffic with this profile. With traffic capture on, Bash commands that use the capture proxy (curl, wget, python-requests) are paced too, by rate and jitter; raw scanners such as nmap and commands that bypass the proxy are not. These flags remain the brake: do not rely on the proxy.",
|
||||
} {
|
||||
if !strings.Contains(got, want) {
|
||||
t.Errorf("%s: missing %q in:\n%s", c.name, want, got)
|
||||
@@ -59,18 +59,23 @@ func TestPromptBuildersNoiseBlock(t *testing.T) {
|
||||
"mainagent": func() string { return mainAgentSystem("goal", "/data", "/work") },
|
||||
"chat": func() string { return chatSystem("auto", "/data", "/work") },
|
||||
}
|
||||
for _, env := range []string{"", "off"} {
|
||||
t.Setenv("ARTEX_NOISE_PROFILE", env)
|
||||
prev := egress.CurrentProfile()
|
||||
t.Cleanup(func() { egress.SetProfile(prev) })
|
||||
// The block follows the current profile, not the environment.
|
||||
t.Setenv("ARTEX_NOISE_PROFILE", "aggressive")
|
||||
egress.SetProfile(egress.Profile{Name: "off"})
|
||||
for name, build := range builders {
|
||||
if got := build(); strings.Contains(got, noiseProfileHeader) {
|
||||
t.Errorf("%s with profile off contains the noise block", name)
|
||||
}
|
||||
}
|
||||
for _, profile := range []string{"quiet", "balanced"} {
|
||||
egress.SetProfile(egress.Profiles[profile])
|
||||
for name, build := range builders {
|
||||
if got := build(); strings.Contains(got, noiseProfileHeader) {
|
||||
t.Errorf("%s with profile %q contains the noise block", name, env)
|
||||
got := build()
|
||||
if !strings.Contains(got, noiseProfileHeader+profile) || strings.Contains(got, noiseProfileHeader+"aggressive") {
|
||||
t.Errorf("%s after SetProfile(%s) is missing its block or shows the env profile", name, profile)
|
||||
}
|
||||
}
|
||||
}
|
||||
t.Setenv("ARTEX_NOISE_PROFILE", "quiet")
|
||||
for name, build := range builders {
|
||||
if got := build(); !strings.Contains(got, noiseProfileHeader+"quiet") {
|
||||
t.Errorf("%s with quiet is missing the noise block", name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+4
-1
@@ -567,9 +567,12 @@ func (p *Planner) Plan(ctx context.Context, taskID int64, as *db.AssetStore, ts
|
||||
if hooks != nil { // typed-nil guard: only set when concrete (avoids harness panic)
|
||||
opts.Hooks = hooks
|
||||
}
|
||||
// The session id is the provider-side identity (session header, usage metering)
|
||||
// with or without persistence; Norma only puts it on ctx when it writes a transcript.
|
||||
opts.SessionID = fmt.Sprintf("exp%d-planner", ts.ID())
|
||||
ctx = transcript.WithSessionID(ctx, opts.SessionID)
|
||||
if p.tx != nil { // persist raw LLM conversation; one accumulating file per task's planner
|
||||
opts.Transcript = p.tx
|
||||
opts.SessionID = fmt.Sprintf("exp%d-planner", ts.ID())
|
||||
}
|
||||
// Experimental: when on, noa takes over context compression (archives persist under <workDir>/noa/<SessionID>).
|
||||
noaSession := fmt.Sprintf("exp%d-planner", ts.ID())
|
||||
|
||||
@@ -49,3 +49,26 @@ func TestProxyEnvInjectsCAWhenRecording(t *testing.T) {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Only the recording proxy (a CA is present) gets the pace marker; the global
|
||||
// proxy of the capture-off path is passed as configured.
|
||||
func TestProxyEnvMarksOnlyTheRecordingProxy(t *testing.T) {
|
||||
proxyVars := map[string]bool{"HTTP_PROXY": true, "HTTPS_PROXY": true, "ALL_PROXY": true}
|
||||
check := func(env []string, want, what string) {
|
||||
t.Helper()
|
||||
n := 0
|
||||
for _, e := range env {
|
||||
if k, v, _ := strings.Cut(e, "="); proxyVars[strings.ToUpper(k)] {
|
||||
n++
|
||||
if v != want {
|
||||
t.Errorf("%s: %s, want %s", what, e, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
if n != 6 {
|
||||
t.Errorf("%s: %d proxy variables, want 6", what, n)
|
||||
}
|
||||
}
|
||||
check(proxyEnv("http://127.0.0.1:8788", "/data/ca.pem"), "http://artex-pace:artex-pace@127.0.0.1:8788", "capture on")
|
||||
check(proxyEnv("http://u:p@10.0.0.1:3128", ""), "http://u:p@10.0.0.1:3128", "capture off")
|
||||
}
|
||||
|
||||
@@ -0,0 +1,19 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/Autumn-27/norma/agentcore"
|
||||
)
|
||||
|
||||
// resumeSession reloads the persisted conversation of session id into s. Norma's
|
||||
// transcript.Store.Load treats a missing file as a fresh session (no records, no
|
||||
// error); any other failure, such as a corrupt line or an unreadable path, is
|
||||
// returned so the caller stops before the first model or tool call instead of
|
||||
// appending a new conversation to the damaged file.
|
||||
func resumeSession(s *agentcore.Session, id string) error {
|
||||
if err := s.Resume(id); err != nil {
|
||||
return fmt.Errorf("restore Norma session %q: %w", id, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,305 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"iter"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
"github.com/Autumn-27/norma/transcript"
|
||||
)
|
||||
|
||||
// sessionProbe is a provider that answers every call with one text turn and
|
||||
// records the session id it saw on ctx plus the request messages.
|
||||
type sessionProbe struct {
|
||||
mu sync.Mutex
|
||||
ids []string
|
||||
msgs [][]llm.Message
|
||||
}
|
||||
|
||||
func (p *sessionProbe) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
|
||||
p.mu.Lock()
|
||||
p.ids = append(p.ids, transcript.SessionIDFrom(ctx))
|
||||
p.msgs = append(p.msgs, req.Messages)
|
||||
p.mu.Unlock()
|
||||
return func(yield func(llm.StreamEvent, error) bool) {
|
||||
if yield(llm.StreamEvent{Type: llm.SETextDelta, Text: "done"}, nil) {
|
||||
yield(llm.StreamEvent{Type: llm.SEMessageDelta, StopReason: "end_turn"}, nil)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (p *sessionProbe) Complete(ctx context.Context, req llm.CompletionRequest) (llm.Message, string, llm.Usage, error) {
|
||||
acc := llm.NewAccumulator()
|
||||
for ev, err := range p.Stream(ctx, req) {
|
||||
if err != nil {
|
||||
return llm.Message{}, "", llm.Usage{}, err
|
||||
}
|
||||
acc.Add(ev)
|
||||
}
|
||||
return acc.Message(), acc.StopReason, acc.Usage, nil
|
||||
}
|
||||
|
||||
func (p *sessionProbe) seen() []string {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
return append([]string(nil), p.ids...)
|
||||
}
|
||||
|
||||
// jsonlFiles lists every .jsonl file under dir.
|
||||
func jsonlFiles(t *testing.T, dir string) []string {
|
||||
t.Helper()
|
||||
var out []string
|
||||
err := filepath.WalkDir(dir, func(path string, d os.DirEntry, err error) error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !d.IsDir() && strings.HasSuffix(path, ".jsonl") {
|
||||
out = append(out, path)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// damageTranscript makes the main transcript of id corrupt (a non-JSON line) or
|
||||
// unreadable (a directory). It returns the bytes a corrupt file must keep.
|
||||
func damageTranscript(t *testing.T, store *transcript.Store, id string, unreadable bool) string {
|
||||
t.Helper()
|
||||
path := store.MainPath(id)
|
||||
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if unreadable {
|
||||
if err := os.Mkdir(path, 0o700); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return ""
|
||||
}
|
||||
const original = "{corrupt transcript\n"
|
||||
if err := os.WriteFile(path, []byte(original), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return original
|
||||
}
|
||||
|
||||
// assertTranscriptUntouched checks that a damaged transcript was neither
|
||||
// rewritten nor appended to.
|
||||
func assertTranscriptUntouched(t *testing.T, store *transcript.Store, id, original string, unreadable bool) {
|
||||
t.Helper()
|
||||
path := store.MainPath(id)
|
||||
if unreadable {
|
||||
entries, err := os.ReadDir(path)
|
||||
if err != nil || len(entries) != 0 {
|
||||
t.Fatalf("unreadable transcript changed: entries=%v err=%v", entries, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
b, err := os.ReadFile(path)
|
||||
if err != nil || string(b) != original {
|
||||
t.Fatalf("damaged transcript changed: %q err=%v", b, err)
|
||||
}
|
||||
}
|
||||
|
||||
func assertRestoreError(t *testing.T, err error, calls int) {
|
||||
t.Helper()
|
||||
if err == nil || !strings.Contains(err.Error(), "restore Norma session") || calls != 0 {
|
||||
t.Fatalf("err=%v provider calls=%d", err, calls)
|
||||
}
|
||||
}
|
||||
|
||||
var damageCases = []struct {
|
||||
name string
|
||||
unreadable bool
|
||||
}{{"corrupt", false}, {"unreadable-directory", true}}
|
||||
|
||||
func TestChatSessionIdentityWithoutTranscript(t *testing.T) {
|
||||
p := &sessionProbe{}
|
||||
workDir := t.TempDir()
|
||||
c := NewChatAgent(p, "test-model", workDir, nil, 100000)
|
||||
for range 2 {
|
||||
if _, err := c.Chat(t.Context(), "assistant", "conv-stable", "hello", 1, 0, false, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if ids := p.seen(); len(ids) != 2 || ids[0] != "conv-stable" || ids[1] != "conv-stable" {
|
||||
t.Fatalf("provider session ids: %v", ids)
|
||||
}
|
||||
if files := jsonlFiles(t, workDir); len(files) != 0 {
|
||||
t.Fatalf("transcript written without persistence: %v", files)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatMissingTranscriptStartsFresh(t *testing.T) {
|
||||
store := transcript.NewStore(t.TempDir())
|
||||
p := &sessionProbe{}
|
||||
c := NewChatAgent(p, "test-model", t.TempDir(), store, 100000)
|
||||
if _, err := c.Chat(t.Context(), "assistant", "conv-new", "hello", 1, 0, false, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if ids := p.seen(); len(ids) != 1 || ids[0] != "conv-new" {
|
||||
t.Fatalf("provider session ids: %v", ids)
|
||||
}
|
||||
recs, err := store.Load(store.MainPath("conv-new"))
|
||||
if err != nil || len(recs) == 0 {
|
||||
t.Fatalf("fresh transcript not written: %d records, err=%v", len(recs), err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatRefusesDamagedTranscript(t *testing.T) {
|
||||
for _, tc := range damageCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
store := transcript.NewStore(t.TempDir())
|
||||
original := damageTranscript(t, store, "conv-broken", tc.unreadable)
|
||||
p := &sessionProbe{}
|
||||
c := NewChatAgent(p, "test-model", t.TempDir(), store, 100000)
|
||||
_, err := c.Chat(t.Context(), "assistant", "conv-broken", "hello", 1, 0, false, nil)
|
||||
assertRestoreError(t, err, len(p.seen()))
|
||||
assertTranscriptUntouched(t, store, "conv-broken", original, tc.unreadable)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatResumesSignedThinkingHistory(t *testing.T) {
|
||||
store := transcript.NewStore(t.TempDir())
|
||||
w := store.NewWriter("conv-signed", "")
|
||||
w.RecordMessage(llm.UserText("earlier question"), llm.Usage{})
|
||||
w.RecordMessage(llm.Message{Role: llm.RoleAssistant, Content: []llm.ContentBlock{
|
||||
{Type: llm.BlockThinking, Thinking: "reasoning", Signature: "signed"},
|
||||
{Type: llm.BlockRedactedThinking, Data: "opaque"},
|
||||
{Type: llm.BlockText, Text: "earlier answer"},
|
||||
}}, llm.Usage{})
|
||||
if err := w.Err(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
p := &sessionProbe{}
|
||||
c := NewChatAgent(p, "test-model", t.TempDir(), store, 100000)
|
||||
if _, err := c.Chat(t.Context(), "assistant", "conv-signed", "next question", 1, 0, false, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
if len(p.msgs) != 1 {
|
||||
t.Fatalf("provider calls: %d", len(p.msgs))
|
||||
}
|
||||
var thinking, redacted bool
|
||||
for _, m := range p.msgs[0] {
|
||||
for _, b := range m.Content {
|
||||
thinking = thinking || (b.Type == llm.BlockThinking && b.Thinking == "reasoning" && b.Signature == "signed")
|
||||
redacted = redacted || (b.Type == llm.BlockRedactedThinking && b.Data == "opaque")
|
||||
}
|
||||
}
|
||||
if !thinking || !redacted {
|
||||
t.Fatalf("resumed history lost signed thinking=%v redacted=%v: %+v", thinking, redacted, p.msgs[0])
|
||||
}
|
||||
}
|
||||
|
||||
// sessionFixture is one exploration with one intent on the test database.
|
||||
type sessionFixture struct {
|
||||
expID int64
|
||||
intentID int64
|
||||
ts *db.ExplorationStore
|
||||
intent *db.Node
|
||||
}
|
||||
|
||||
func newSessionFixture(t *testing.T) sessionFixture {
|
||||
t.Helper()
|
||||
d := testDB(t)
|
||||
t.Cleanup(func() { _ = d.Close() })
|
||||
expID, err := d.CreateExploration("session identity probe", "keep the session id stable")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _, _ = d.Exec(`DELETE FROM explorations WHERE id=$1`, expID) })
|
||||
ts := d.Exploration(expID)
|
||||
payload := map[string]any{"summary": "session probe"}
|
||||
intentID, err := ts.AddNode("intent", payload, 0, "running", "planner", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
raw, _ := json.Marshal(payload)
|
||||
return sessionFixture{expID: expID, intentID: intentID, ts: ts, intent: &db.Node{ID: intentID, Payload: raw}}
|
||||
}
|
||||
|
||||
// run drives one agent kind once and returns the session id it should use.
|
||||
func (f sessionFixture) run(t *testing.T, kind string, p *sessionProbe, tx *transcript.Store) (string, error) {
|
||||
t.Helper()
|
||||
workDir := t.TempDir()
|
||||
switch kind {
|
||||
case "worker":
|
||||
_, _, err := NewWorker(p, "test-model", workDir, tx, 0, 3).Execute(t.Context(), "work#1", f.expID, nil, f.ts, f.intent, passHooks{}, nil, nil, nil)
|
||||
return WorkerSessionID(f.expID, f.intentID), err
|
||||
case "planner":
|
||||
_, _, err := NewPlanner(p, "test-model", workDir, tx, 0, 3).Plan(t.Context(), f.expID, nil, f.ts, "probe goal", nil, passHooks{}, nil)
|
||||
return fmt.Sprintf("exp%d-planner", f.expID), err
|
||||
case "mainagent":
|
||||
_, err := NewMainAgent(p, "test-model", workDir, tx, 0, 3).Chat(t.Context(), f.expID, 0, nil, f.ts, "probe goal", "status?", passHooks{}, nil, nil, nil, nil, nil)
|
||||
return fmt.Sprintf("exp%d-main", f.expID), err
|
||||
case "mainagent-seg":
|
||||
_, err := NewMainAgent(p, "test-model", workDir, tx, 0, 3).Chat(t.Context(), f.expID, 2, nil, f.ts, "probe goal", "status?", passHooks{}, nil, nil, nil, nil, nil)
|
||||
return fmt.Sprintf("exp%d-main-s2", f.expID), err
|
||||
}
|
||||
t.Fatalf("unknown agent kind %q", kind)
|
||||
return "", nil
|
||||
}
|
||||
|
||||
func TestAgentSessionIdentityWithAndWithoutTranscript(t *testing.T) {
|
||||
f := newSessionFixture(t)
|
||||
for _, kind := range []string{"worker", "planner", "mainagent", "mainagent-seg"} {
|
||||
t.Run(kind, func(t *testing.T) {
|
||||
bare := &sessionProbe{}
|
||||
want, err := f.run(t, kind, bare, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
store := transcript.NewStore(t.TempDir())
|
||||
persisted := &sessionProbe{}
|
||||
if _, err := f.run(t, kind, persisted, store); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, ids := range [][]string{bare.seen(), persisted.seen()} {
|
||||
if len(ids) == 0 {
|
||||
t.Fatal("provider never called")
|
||||
}
|
||||
for _, id := range ids {
|
||||
if id != want {
|
||||
t.Fatalf("provider session ids %v, want %q", ids, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
if recs, err := store.Load(store.MainPath(want)); err != nil || len(recs) == 0 {
|
||||
t.Fatalf("transcript %q not written under the session id: %d records, err=%v", want, len(recs), err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentRefusesDamagedTranscript(t *testing.T) {
|
||||
f := newSessionFixture(t)
|
||||
for _, kind := range []string{"worker", "mainagent"} {
|
||||
for _, tc := range damageCases {
|
||||
t.Run(kind+"/"+tc.name, func(t *testing.T) {
|
||||
store := transcript.NewStore(t.TempDir())
|
||||
id := WorkerSessionID(f.expID, f.intentID)
|
||||
if kind == "mainagent" {
|
||||
id = fmt.Sprintf("exp%d-main", f.expID)
|
||||
}
|
||||
original := damageTranscript(t, store, id, tc.unreadable)
|
||||
p := &sessionProbe{}
|
||||
_, err := f.run(t, kind, p, store)
|
||||
assertRestoreError(t, err, len(p.seen()))
|
||||
assertTranscriptUntouched(t, store, id, original, tc.unreadable)
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
+4
-4
@@ -46,9 +46,9 @@ const stagnationMinRepeats = 3
|
||||
|
||||
func detectStagnation(done, findings []*db.Node) []map[string]any {
|
||||
type dir struct {
|
||||
firstID int64
|
||||
repeats int
|
||||
summary string
|
||||
firstID int64
|
||||
repeats int
|
||||
summary string
|
||||
}
|
||||
dirs := map[string]*dir{}
|
||||
order := []string{}
|
||||
@@ -551,7 +551,7 @@ func (t *ToolSet) graphOverviewData() map[string]any {
|
||||
// with the data so no prompt-template change is needed on deployed installs.
|
||||
if stale := detectStagnation(recentDone, vulnNodes); len(stale) > 0 {
|
||||
out["stagnation_warning"] = map[string]any{
|
||||
"directions": stale,
|
||||
"directions": stale,
|
||||
"instruction": "these directions were dispatched repeatedly with zero findings produced since their first attempt — do NOT re-dispatch the same approach again; open a materially different direction, or judge the goal on the existing evidence instead of grinding",
|
||||
}
|
||||
}
|
||||
|
||||
+5
-5
@@ -27,11 +27,12 @@ func withWebFetch(ctx context.Context, tools []actool.CoreTool, proxy, caCert st
|
||||
gated := scoped && scope.Enabled != nil && scope.Enabled()
|
||||
webFetch := actool.NewWebFetch(actool.WebFetchConfig{Proxy: proxy, CACert: caCert, CheckRedirect: intercept.RedirectCheck("WebFetch"),
|
||||
DialContext: intercept.PinnedDialer("WebFetch", scoped), DisableDirectRetry: gated})
|
||||
return append(tools[:len(tools):len(tools)], governedFetch{webFetch, egressGovernor()})
|
||||
return append(tools[:len(tools):len(tools)], governedFetch{webFetch})
|
||||
}
|
||||
|
||||
// egressGovernor is the governor of WebFetch; a test may pin it.
|
||||
var egressGovernor = egress.Shared
|
||||
// egressGovernor resolves the governor of a WebFetch call, once per call so a
|
||||
// changed noise profile applies to the next call; a test may pin it.
|
||||
var egressGovernor = egress.Current
|
||||
|
||||
// governedFetch runs each WebFetch call as one send under the egress governor
|
||||
// (off by default; then the call runs unchanged). Under a noise profile it
|
||||
@@ -42,7 +43,6 @@ var egressGovernor = egress.Shared
|
||||
// 429/503 is not retried.
|
||||
type governedFetch struct {
|
||||
actool.CoreTool
|
||||
g *egress.Governor
|
||||
}
|
||||
|
||||
func (f governedFetch) Call(ctx context.Context, input json.RawMessage, tc *actool.ToolContext) (actool.Result, error) {
|
||||
@@ -56,7 +56,7 @@ func (f governedFetch) Call(ctx context.Context, input json.RawMessage, tc *acto
|
||||
if err != nil || u.Host == "" {
|
||||
return f.CoreTool.Call(ctx, input, tc) // the tool reports the bad URL
|
||||
}
|
||||
return egress.Do(f.g, ctx, u, func() (actool.Result, error) {
|
||||
return egress.Do(egressGovernor(), ctx, u, func() (actool.Result, error) {
|
||||
return f.CoreTool.Call(ctx, input, tc)
|
||||
})
|
||||
}
|
||||
|
||||
@@ -27,7 +27,11 @@ func TestWebFetchProfileDoesNotCache(t *testing.T) {
|
||||
checks.Add(1)
|
||||
return nil
|
||||
}})
|
||||
tool := governedFetch{inner, egress.New(egress.Profiles["balanced"])}
|
||||
balanced := egress.New(egress.Profiles["balanced"])
|
||||
old := egressGovernor
|
||||
egressGovernor = func() *egress.Governor { return balanced }
|
||||
t.Cleanup(func() { egressGovernor = old })
|
||||
tool := governedFetch{inner}
|
||||
ctx := intercept.WithTargetScope(t.Context(), intercept.TargetScope{TaskID: "7"})
|
||||
for range 2 {
|
||||
res, err := tool.Call(ctx, json.RawMessage(fetchInputJSON(pair.start)), &actool.ToolContext{})
|
||||
|
||||
+20
-2
@@ -11,6 +11,7 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/egress"
|
||||
"github.com/Autumn-27/artex/intercept"
|
||||
"github.com/Autumn-27/norma/agentcore"
|
||||
"github.com/Autumn-27/norma/harness"
|
||||
@@ -212,10 +213,19 @@ func (w *Worker) SetWebSearch(o WebSearchOpts) { w.webSearch = o }
|
||||
// honor the proxy vars. ALL_PROXY is set too so a socks5 egress proxy (which curl
|
||||
// only reads from ALL_PROXY, not HTTP(S)_PROXY) works in the capture-off path.
|
||||
// Empty proxyAddr → nil (direct, unchanged env).
|
||||
//
|
||||
// With a MITM CA (capture on) the proxy URL carries the pace marker userinfo
|
||||
// (egress.MarkProxyURL): the recording proxy then paces this connection's
|
||||
// requests by the noise profile. Only Bash gets the marker; WebFetch, the http
|
||||
// custom tools and enrich keep the plain URL, since egress paces them already.
|
||||
// The global proxy (capture off) is left as configured.
|
||||
func proxyEnv(proxyAddr, caCert string) []string {
|
||||
if proxyAddr == "" {
|
||||
return nil
|
||||
}
|
||||
if caCert != "" {
|
||||
proxyAddr = egress.MarkProxyURL(proxyAddr)
|
||||
}
|
||||
env := []string{
|
||||
"HTTP_PROXY=" + proxyAddr, "HTTPS_PROXY=" + proxyAddr,
|
||||
"http_proxy=" + proxyAddr, "https_proxy=" + proxyAddr,
|
||||
@@ -530,9 +540,12 @@ func (w *Worker) execute(ctx context.Context, name string, taskID int64, as *db.
|
||||
if hooks != nil { // typed-nil guard: only set when concrete (avoids harness panic)
|
||||
opts.Hooks = hooks
|
||||
}
|
||||
// The session id is the provider-side identity (session header, usage metering)
|
||||
// with or without persistence; Norma only puts it on ctx when it writes a transcript.
|
||||
opts.SessionID = WorkerSessionID(ts.ID(), intent.ID)
|
||||
ctx = transcript.WithSessionID(ctx, opts.SessionID)
|
||||
if w.tx != nil { // persist raw LLM conversation; one file per worked intent
|
||||
opts.Transcript = w.tx
|
||||
opts.SessionID = WorkerSessionID(ts.ID(), intent.ID)
|
||||
}
|
||||
intentID := intent.ID
|
||||
emitWrap := func(r db.Activity) {
|
||||
@@ -564,7 +577,12 @@ func (w *Worker) execute(ctx context.Context, name string, taskID int64, as *db.
|
||||
// restarting from scratch.
|
||||
alreadyRecorded := false
|
||||
if w.tx != nil {
|
||||
_ = s.Resume(opts.SessionID)
|
||||
// A corrupt or unreadable transcript stops the run before any model or tool
|
||||
// call. No terminal reason: the engine marks the intent blocked without the
|
||||
// model_error reruns, which would only fail the same way again.
|
||||
if err := resumeSession(s, opts.SessionID); err != nil {
|
||||
return "", tsx.Writes(), err
|
||||
}
|
||||
alreadyRecorded = requestID != "" && hasWorkerChatMessage(s.Messages(), requestID)
|
||||
if len(s.Messages()) > 0 && message == "" {
|
||||
seedUnlockFromHistory(s.Messages(), def.UnlockSkill)
|
||||
|
||||
@@ -7,12 +7,12 @@ import "time"
|
||||
// these match the *target asset*: an exact/fuzzy domain·ip·url, or a CIDR range.
|
||||
// This layer only stores rules; the matching/enforcement logic lives elsewhere.
|
||||
type AssetInterceptRule struct {
|
||||
ID int64 `json:"id"`
|
||||
Enabled bool `json:"enabled"`
|
||||
Kind string `json:"kind"` // exact_domain|exact_ip|exact_url|fuzzy_domain|fuzzy_ip|fuzzy_url|cidr
|
||||
Pattern string `json:"pattern"`
|
||||
Note string `json:"note"`
|
||||
Builtin bool `json:"builtin"`
|
||||
ID int64 `json:"id"`
|
||||
Enabled bool `json:"enabled"`
|
||||
Kind string `json:"kind"` // exact_domain|exact_ip|exact_url|fuzzy_domain|fuzzy_ip|fuzzy_url|cidr
|
||||
Pattern string `json:"pattern"`
|
||||
Note string `json:"note"`
|
||||
Builtin bool `json:"builtin"`
|
||||
// Action is only used by task-level rules: 'block'=intercept, 'allow'=allow (whitelist).
|
||||
// Global rules (asset_intercept_rules) do not have this column; it is always empty and treated as intercept.
|
||||
Action string `json:"action,omitempty"`
|
||||
|
||||
+6
-6
@@ -1168,11 +1168,11 @@ func TestQueryDSLInScopeMembership(t *testing.T) {
|
||||
}
|
||||
|
||||
stamp := time.Now().UnixNano()
|
||||
root := fmt.Sprintf("sc%d.invalid", stamp) // in-scope root domain (task)
|
||||
srcRoot := fmt.Sprintf("src%d.invalid", stamp) // in-scope via source task
|
||||
out := fmt.Sprintf("out%d.invalid", stamp) // out of every scope
|
||||
marker := fmt.Sprintf("mk%d", stamp) // bare-text token present in all rows
|
||||
foreignTask := stamp + 777 // asset produced by an unrelated task
|
||||
root := fmt.Sprintf("sc%d.invalid", stamp) // in-scope root domain (task)
|
||||
srcRoot := fmt.Sprintf("src%d.invalid", stamp) // in-scope via source task
|
||||
out := fmt.Sprintf("out%d.invalid", stamp) // out of every scope
|
||||
marker := fmt.Sprintf("mk%d", stamp) // bare-text token present in all rows
|
||||
foreignTask := stamp + 777 // asset produced by an unrelated task
|
||||
|
||||
// Scope: task owns root; source task owns srcRoot; task owns an IP /24.
|
||||
for _, sc := range []struct {
|
||||
@@ -1200,7 +1200,7 @@ func TestQueryDSLInScopeMembership(t *testing.T) {
|
||||
{"subdomain", "api." + root, root, "", "", "", marker}, // under task root
|
||||
{"service", "www." + root, root, "https://www." + root + "/", "", "", marker}, // service under task root
|
||||
{"endpoint", "www." + root, root, "https://www." + root + "/a?" + marker + "=1", "GET", "", ""},
|
||||
{"subdomain", "dev." + srcRoot, srcRoot, "", "", "", marker}, // under source-task root
|
||||
{"subdomain", "dev." + srcRoot, srcRoot, "", "", "", marker}, // under source-task root
|
||||
{"endpoint", "198.51.100.9", "198.51.100.9", "http://198.51.100.9:8080/" + marker, "GET", "", ""}, // IP-literal host, ip col empty
|
||||
{"subdomain", "x." + out, out, "", "", "", marker}, // out of scope
|
||||
}
|
||||
|
||||
+50
-26
@@ -1486,11 +1486,13 @@ type TaskListMetrics struct {
|
||||
Findings FindingSeverityCounts // this task's vulnerability count in the findings table (bucketed by severity)
|
||||
}
|
||||
|
||||
// TaskListMetricsAll returns list aggregates for every live task in one query.
|
||||
// The lateral lookups use the per-exploration indexes instead of grouping the
|
||||
// complete activity history on every poll.
|
||||
func (d *DB) TaskListMetricsAll() (map[int64]TaskListMetrics, error) {
|
||||
rows, err := d.Query(`
|
||||
// taskListMetricsQuery is the select shared by TaskListMetricsAll and
|
||||
// TaskListMetrics; each appends its own WHERE on task. The lateral lookups use
|
||||
// the per-exploration indexes instead of grouping the complete activity history
|
||||
// on every poll. latest_activity orders by the column activity.created_at, so
|
||||
// idx_act_latest (exploration_id, created_at DESC) serves the LIMIT 1; ordering
|
||||
// by the output alias would sort every activity row of the task.
|
||||
const taskListMetricsQuery = `
|
||||
SELECT task.exploration_id,
|
||||
COALESCE(token_metrics.input_tokens,0),
|
||||
COALESCE(token_metrics.output_tokens,0),
|
||||
@@ -1514,10 +1516,10 @@ func (d *DB) TaskListMetricsAll() (map[int64]TaskListMetrics, error) {
|
||||
WHERE exploration_id=task.exploration_id AND kind='result'
|
||||
) token_metrics ON true
|
||||
LEFT JOIN LATERAL (
|
||||
SELECT EXTRACT(EPOCH FROM created_at)::bigint AS created_at
|
||||
SELECT EXTRACT(EPOCH FROM activity.created_at)::bigint AS created_at
|
||||
FROM activity
|
||||
WHERE exploration_id=task.exploration_id
|
||||
ORDER BY created_at DESC
|
||||
WHERE activity.exploration_id=task.exploration_id
|
||||
ORDER BY activity.created_at DESC
|
||||
LIMIT 1
|
||||
) latest_activity ON true
|
||||
LEFT JOIN LATERAL (
|
||||
@@ -1539,30 +1541,41 @@ func (d *DB) TaskListMetricsAll() (map[int64]TaskListMetrics, error) {
|
||||
FROM findings
|
||||
WHERE task_id=task.id
|
||||
) finding_metrics ON true
|
||||
WHERE task.deleted_at IS NULL`)
|
||||
WHERE task.deleted_at IS NULL`
|
||||
|
||||
// scanTaskListMetrics reads one row of taskListMetricsQuery.
|
||||
func scanTaskListMetrics(row interface{ Scan(...any) error }) (int64, TaskListMetrics, error) {
|
||||
var explorationID int64
|
||||
var metrics TaskListMetrics
|
||||
err := row.Scan(
|
||||
&explorationID,
|
||||
&metrics.Tokens.InputTokens,
|
||||
&metrics.Tokens.OutputTokens,
|
||||
&metrics.Tokens.CacheReadTokens,
|
||||
&metrics.Tokens.CacheWriteTokens,
|
||||
&metrics.LastActivity,
|
||||
&metrics.Goals.Total,
|
||||
&metrics.Goals.Met,
|
||||
&metrics.RunningIntents,
|
||||
&metrics.Findings.Critical,
|
||||
&metrics.Findings.High,
|
||||
&metrics.Findings.Medium,
|
||||
&metrics.Findings.Low,
|
||||
)
|
||||
return explorationID, metrics, err
|
||||
}
|
||||
|
||||
// TaskListMetricsAll returns list aggregates for every live task in one query.
|
||||
func (d *DB) TaskListMetricsAll() (map[int64]TaskListMetrics, error) {
|
||||
rows, err := d.Query(taskListMetricsQuery)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
out := map[int64]TaskListMetrics{}
|
||||
for rows.Next() {
|
||||
var explorationID int64
|
||||
var metrics TaskListMetrics
|
||||
if err := rows.Scan(
|
||||
&explorationID,
|
||||
&metrics.Tokens.InputTokens,
|
||||
&metrics.Tokens.OutputTokens,
|
||||
&metrics.Tokens.CacheReadTokens,
|
||||
&metrics.Tokens.CacheWriteTokens,
|
||||
&metrics.LastActivity,
|
||||
&metrics.Goals.Total,
|
||||
&metrics.Goals.Met,
|
||||
&metrics.RunningIntents,
|
||||
&metrics.Findings.Critical,
|
||||
&metrics.Findings.High,
|
||||
&metrics.Findings.Medium,
|
||||
&metrics.Findings.Low,
|
||||
); err != nil {
|
||||
explorationID, metrics, err := scanTaskListMetrics(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out[explorationID] = metrics
|
||||
@@ -1570,6 +1583,17 @@ func (d *DB) TaskListMetricsAll() (map[int64]TaskListMetrics, error) {
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// TaskListMetrics returns the same aggregates for the one live task of
|
||||
// exploration expID, with the same laterals as TaskListMetricsAll. A missing or
|
||||
// deleted task gives zero metrics, as a missing key of TaskListMetricsAll does.
|
||||
func (d *DB) TaskListMetrics(expID int64) (TaskListMetrics, error) {
|
||||
_, metrics, err := scanTaskListMetrics(d.QueryRow(taskListMetricsQuery+` AND task.exploration_id=$1`, expID))
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return TaskListMetrics{}, nil
|
||||
}
|
||||
return metrics, err
|
||||
}
|
||||
|
||||
// GoalCountsAll returns goal totals (total / met) per exploration id, one query
|
||||
// for all tasks — used by listTasks so the task list shows progress for every task.
|
||||
func (d *DB) GoalCountsAll() (map[int64]GoalCounts, error) {
|
||||
|
||||
@@ -117,3 +117,133 @@ VALUES ($1,'goal','{}','met'), ($1,'goal','{}','open')`, task.ExplorationID); er
|
||||
t.Fatalf("unexpected goal metrics: %+v", metrics.Goals)
|
||||
}
|
||||
}
|
||||
|
||||
// TaskListMetrics(expID) is exactly the row TaskListMetricsAll returns for that
|
||||
// exploration, on every counter, and TaskListMetricsAll still lists every task.
|
||||
func TestTaskListMetricsMatchesAll(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
newest := time.Now().Add(-time.Minute).Truncate(time.Second)
|
||||
seed := func(name string, n int) *Task {
|
||||
t.Helper()
|
||||
task, err := d.CreateTask(name, "goal", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_, err := d.DeleteTaskCascade(task.ID, false, true)
|
||||
cleanupErr(t, "task "+name, err)
|
||||
})
|
||||
exp := task.ExplorationID
|
||||
// n result rows with tokens; the newest activity first, then an older row
|
||||
// with a higher id, so ordering by id instead of created_at would fail.
|
||||
for _, q := range []struct {
|
||||
sql string
|
||||
args []any
|
||||
}{
|
||||
{`INSERT INTO activity(exploration_id,kind,input_tokens,output_tokens,cache_read_tokens,cache_write_tokens,created_at)
|
||||
SELECT $1,'result',10,5,3,1,$3::timestamptz - interval '1 hour' FROM generate_series(1,$2::int)`, []any{exp, n, newest}},
|
||||
{`INSERT INTO activity(exploration_id,kind,created_at) VALUES ($1,'text',$2::timestamptz)`, []any{exp, newest}},
|
||||
{`INSERT INTO activity(exploration_id,kind,created_at) SELECT $1,'text',$3::timestamptz - interval '2 hours' FROM generate_series(1,$2::int)`, []any{exp, n, newest}},
|
||||
{`INSERT INTO exploration_nodes(exploration_id,kind,payload,state)
|
||||
SELECT $1,'goal','{}',CASE WHEN g<=$2::int THEN 'met' ELSE 'open' END FROM generate_series(1,$2::int+1) g`, []any{exp, n}},
|
||||
{`INSERT INTO exploration_nodes(exploration_id,kind,payload,state)
|
||||
SELECT $1,'intent','{}',CASE WHEN g<=$2::int THEN 'running' ELSE 'done' END FROM generate_series(1,$2::int+2) g`, []any{exp, n}},
|
||||
} {
|
||||
if _, err := d.Exec(q.sql, q.args...); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if _, err := d.Exec(`INSERT INTO findings(task_id,severity)
|
||||
SELECT $1, s FROM unnest(ARRAY['critical','high','high','medium','low','info']) s, generate_series(1,$2)`, task.ID, n); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return task
|
||||
}
|
||||
a := seed("task-list-metrics-a", 2)
|
||||
b := seed("task-list-metrics-b", 3)
|
||||
|
||||
all, err := d.TaskListMetricsAll()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, c := range []struct {
|
||||
task *Task
|
||||
n int
|
||||
}{{a, 2}, {b, 3}} {
|
||||
row, ok := all[c.task.ExplorationID]
|
||||
if !ok {
|
||||
t.Fatalf("TaskListMetricsAll misses task %d", c.task.ID)
|
||||
}
|
||||
one, err := d.TaskListMetrics(c.task.ExplorationID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if one != row {
|
||||
t.Fatalf("task %d: TaskListMetrics=%+v, TaskListMetricsAll=%+v", c.task.ID, one, row)
|
||||
}
|
||||
n := c.n
|
||||
want := TaskListMetrics{
|
||||
Tokens: TokenUsage{InputTokens: 10 * n, OutputTokens: 5 * n, CacheReadTokens: 3 * n, CacheWriteTokens: n},
|
||||
LastActivity: newest.Unix(),
|
||||
Goals: GoalCounts{Total: n + 1, Met: n},
|
||||
RunningIntents: n,
|
||||
Findings: FindingSeverityCounts{Critical: n, High: 2 * n, Medium: n, Low: n},
|
||||
}
|
||||
if one != want {
|
||||
t.Fatalf("task %d metrics=%+v, want %+v", c.task.ID, one, want)
|
||||
}
|
||||
}
|
||||
if none, err := d.TaskListMetrics(-1); err != nil || none != (TaskListMetrics{}) {
|
||||
t.Fatalf("missing exploration: %+v, %v; want zero metrics", none, err)
|
||||
}
|
||||
}
|
||||
|
||||
// BenchmarkTaskListMetrics compares the per-exploration metrics query with the
|
||||
// all-tasks one on three tasks of 40,000 activity rows each. No time assertion.
|
||||
func BenchmarkTaskListMetrics(b *testing.B) {
|
||||
dsn, _, err := DSN()
|
||||
if err != nil {
|
||||
b.Skipf("no database config (%v) - skipping", err)
|
||||
}
|
||||
d, err := Open(dsn)
|
||||
if err != nil {
|
||||
b.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
b.Cleanup(func() { d.Close() })
|
||||
var exp int64
|
||||
for i := range 3 {
|
||||
task, err := d.CreateTask(fmt.Sprintf("bench-task-list-metrics-%d", i), "goal", nil, 0, 0)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
b.Cleanup(func() { _, _ = d.DeleteTaskCascade(task.ID, false, true) })
|
||||
if _, err := d.Exec(`INSERT INTO activity(exploration_id,kind,input_tokens,created_at)
|
||||
SELECT $1, CASE WHEN g%20=0 THEN 'result' ELSE 'tool_use' END, 100, now() - (g||' seconds')::interval
|
||||
FROM generate_series(1,40000) g`, task.ExplorationID); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
exp = task.ExplorationID
|
||||
}
|
||||
if _, err := d.Exec(`ANALYZE activity`); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
b.Run("one", func(b *testing.B) {
|
||||
for range b.N {
|
||||
if _, err := d.TaskListMetrics(exp); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
})
|
||||
b.Run("all", func(b *testing.B) {
|
||||
for range b.N {
|
||||
if _, err := d.TaskListMetricsAll(); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
@@ -6,14 +6,14 @@ import "testing"
|
||||
// ip/cidr scope path so a target like "10.0.188.136:3000" no longer 404s.
|
||||
func TestStripHostPort(t *testing.T) {
|
||||
cases := []struct{ in, want string }{
|
||||
{"10.0.188.136:3000", "10.0.188.136"}, // the reported IP case
|
||||
{"10.0.188.136", "10.0.188.136"}, // bare IPv4 unchanged
|
||||
{"[2001:db8::1]:8080", "2001:db8::1"}, // bracketed IPv6 + port
|
||||
{"2001:db8::1", "2001:db8::1"}, // bare IPv6 unchanged (has colons)
|
||||
{" 1.2.3.4:80 ", "1.2.3.4"}, // trims surrounding space
|
||||
{"example.com:443", "example.com"}, // domain + port → bare host
|
||||
{"api.example.com:8080", "api.example.com"}, // subdomain + port
|
||||
{"example.com", "example.com"}, // bare domain unchanged
|
||||
{"10.0.188.136:3000", "10.0.188.136"}, // the reported IP case
|
||||
{"10.0.188.136", "10.0.188.136"}, // bare IPv4 unchanged
|
||||
{"[2001:db8::1]:8080", "2001:db8::1"}, // bracketed IPv6 + port
|
||||
{"2001:db8::1", "2001:db8::1"}, // bare IPv6 unchanged (has colons)
|
||||
{" 1.2.3.4:80 ", "1.2.3.4"}, // trims surrounding space
|
||||
{"example.com:443", "example.com"}, // domain + port → bare host
|
||||
{"api.example.com:8080", "api.example.com"}, // subdomain + port
|
||||
{"example.com", "example.com"}, // bare domain unchanged
|
||||
}
|
||||
for _, c := range cases {
|
||||
if got := stripHostPort(c.in); got != c.want {
|
||||
|
||||
+11
-11
@@ -28,19 +28,19 @@ type Task struct {
|
||||
LLMProfileID *int64 `json:"llm_profile_id,omitempty"`
|
||||
// Task-level ordered LLM chain. LLMProfileID remains the compatibility alias
|
||||
// for ActiveLLMProfileID while older API clients still send one profile id.
|
||||
LLMProfileIDs []int64 `json:"llm_profile_ids,omitempty"`
|
||||
ActiveLLMProfileID *int64 `json:"active_llm_profile_id,omitempty"`
|
||||
LLMChainRevision int64 `json:"-"`
|
||||
LLMFailoverState string `json:"llm_failover_state,omitempty"`
|
||||
LLMFailoverReason string `json:"llm_failover_reason,omitempty"`
|
||||
LLMProfileIDs []int64 `json:"llm_profile_ids,omitempty"`
|
||||
ActiveLLMProfileID *int64 `json:"active_llm_profile_id,omitempty"`
|
||||
LLMChainRevision int64 `json:"-"`
|
||||
LLMFailoverState string `json:"llm_failover_state,omitempty"`
|
||||
LLMFailoverReason string `json:"llm_failover_reason,omitempty"`
|
||||
// LLMRotate: per-intent chain-head rotation switch (#85); false = classic
|
||||
// ordered failover from the active head.
|
||||
LLMRotate bool `json:"llm_rotate"`
|
||||
SourceTaskIDs []int64 `json:"source_task_ids,omitempty"`
|
||||
CompanyIDs []int64 `json:"company_ids,omitempty"`
|
||||
ParentRef string `json:"parent_ref,omitempty"` // parent task id (recorded by orchestration spawn; empty=top level)
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
CompletedAt *time.Time `json:"completed_at,omitempty"` // when it entered a terminal state (done/failed/timeout); nil for non-terminal
|
||||
LLMRotate bool `json:"llm_rotate"`
|
||||
SourceTaskIDs []int64 `json:"source_task_ids,omitempty"`
|
||||
CompanyIDs []int64 `json:"company_ids,omitempty"`
|
||||
ParentRef string `json:"parent_ref,omitempty"` // parent task id (recorded by orchestration spawn; empty=top level)
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
CompletedAt *time.Time `json:"completed_at,omitempty"` // when it entered a terminal state (done/failed/timeout); nil for non-terminal
|
||||
// Task-level timeout (see the task-level timeout and wrap-up design doc).
|
||||
TimeoutSeconds int `json:"timeout_seconds"` // 0=no time limit
|
||||
FirstRunAt *time.Time `json:"first_run_at,omitempty"` // when it first actually started running (not created_at); nil=not yet run
|
||||
|
||||
@@ -0,0 +1,179 @@
|
||||
package egress
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"runtime"
|
||||
"sync"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// keepProfile restores the process profile after a test that changes it.
|
||||
func keepProfile(t *testing.T) {
|
||||
prev := CurrentProfile()
|
||||
t.Cleanup(func() { SetProfile(prev) })
|
||||
}
|
||||
|
||||
func TestSetProfileConcurrentAndConsistent(t *testing.T) {
|
||||
keepProfile(t)
|
||||
names := []string{"off", "quiet", "balanced", "aggressive"}
|
||||
var wg sync.WaitGroup
|
||||
for i := range 8 {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for j := range 200 {
|
||||
p, _ := ProfileByName(names[(i+j)%len(names)])
|
||||
SetProfile(p)
|
||||
if g := Current(); g == nil || g.Profile().Name == "" {
|
||||
t.Errorf("Current() = %v during swaps", g)
|
||||
return
|
||||
}
|
||||
_ = CurrentProfile()
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
for _, name := range names {
|
||||
p, _ := ProfileByName(name)
|
||||
g := SetProfile(p)
|
||||
if Current() != g || CurrentProfile().Name != name || Current().Profile() != CurrentProfile() {
|
||||
t.Fatalf("after SetProfile(%s): Current=%p want %p, CurrentProfile=%q", name, Current(), g, CurrentProfile().Name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// A request in flight keeps the governor it started on: its release frees the
|
||||
// old governor's slots, and the new governor starts with no counter.
|
||||
func TestSetProfileInFlightFinishesOnOldGovernor(t *testing.T) {
|
||||
keepProfile(t)
|
||||
old := SetProfile(open)
|
||||
release, err := old.acquire(t.Context(), "a.example")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
fresh := SetProfile(Profiles["balanced"])
|
||||
if Current() != fresh || len(fresh.Snapshot()) != 0 {
|
||||
t.Fatalf("new governor: current=%p fresh=%p snapshot=%v", Current(), fresh, fresh.Snapshot())
|
||||
}
|
||||
if s := old.Snapshot(); len(s) != 1 || s[0].Inflight != 1 {
|
||||
t.Fatalf("old governor before release: %+v", s)
|
||||
}
|
||||
release()
|
||||
if s := old.Snapshot(); len(s) != 1 || s[0].Inflight != 0 || s[0].RPM != 1 || s[0].PeakInflight != 1 {
|
||||
t.Fatalf("old governor after release: %+v", s)
|
||||
}
|
||||
}
|
||||
|
||||
type recordRT struct{ got *http.Request }
|
||||
|
||||
func (r *recordRT) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
r.got = req
|
||||
return &http.Response{StatusCode: 204, Body: http.NoBody, Request: req}, nil
|
||||
}
|
||||
|
||||
// Off stays a strict passthrough, through Transport and Dynamic alike.
|
||||
func TestOffIsPassthroughAfterSetProfile(t *testing.T) {
|
||||
keepProfile(t)
|
||||
t.Setenv("ARTEX_TARGET_USER_AGENT", "artex-test/1")
|
||||
off, ok := ProfileByName("OFF")
|
||||
if !ok || off.Name != "off" || off.UserAgent != "" {
|
||||
t.Fatalf("ProfileByName(OFF) = %+v %v", off, ok)
|
||||
}
|
||||
SetProfile(off)
|
||||
base := &recordRT{}
|
||||
if tr := Current().Transport("x", base); tr != http.RoundTripper(base) {
|
||||
t.Fatalf("off Transport = %T, want the base itself", tr)
|
||||
}
|
||||
req := httptest.NewRequest(http.MethodGet, "http://a.example/", nil)
|
||||
if _, err := Dynamic("x", base).RoundTrip(req); err != nil || base.got != req || req.Header.Get("User-Agent") != "" {
|
||||
t.Fatalf("off Dynamic: err=%v same request=%v UA=%q", err, base.got == req, req.Header.Get("User-Agent"))
|
||||
}
|
||||
if len(Current().Snapshot()) != 0 {
|
||||
t.Fatalf("off counted a request: %v", Current().Snapshot())
|
||||
}
|
||||
// Any other profile carries the user agent.
|
||||
if p, _ := ProfileByName("balanced"); p.UserAgent != "artex-test/1" {
|
||||
t.Fatalf("balanced UserAgent = %q", p.UserAgent)
|
||||
}
|
||||
if _, ok := ProfileByName(""); ok {
|
||||
t.Fatal(`ProfileByName("") accepted`)
|
||||
}
|
||||
if _, ok := ProfileByName("loud"); ok {
|
||||
t.Fatal(`ProfileByName("loud") accepted`)
|
||||
}
|
||||
}
|
||||
|
||||
// Dynamic resolves the governor per request: a client built before
|
||||
// SetProfile is paced by the new profile.
|
||||
func TestDynamicFollowsSetProfile(t *testing.T) {
|
||||
keepProfile(t)
|
||||
srv := newServer(t, 0, nil)
|
||||
SetProfile(Profile{Name: "off"})
|
||||
client := &http.Client{Transport: Dynamic("ns", nil)}
|
||||
get := func() {
|
||||
t.Helper()
|
||||
resp, err := client.Get(srv.URL)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
resp.Body.Close()
|
||||
}
|
||||
get()
|
||||
g := SetProfile(open)
|
||||
get()
|
||||
get()
|
||||
if s := g.Snapshot(); len(s) != 1 || s[0].RPM != 2 {
|
||||
t.Fatalf("snapshot after switch: %+v", s)
|
||||
}
|
||||
}
|
||||
|
||||
// Snapshot reads the counters without starting a new minute; report does.
|
||||
func TestSnapshotDoesNotReset(t *testing.T) {
|
||||
g := New(open)
|
||||
srv := newServer(t, 0, nil)
|
||||
tr := g.Transport("ns", nil)
|
||||
for range 3 {
|
||||
req, _ := http.NewRequest(http.MethodGet, srv.URL, nil)
|
||||
resp, err := tr.RoundTrip(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
resp.Body.Close()
|
||||
}
|
||||
want := HostNoise{Host: HostKey(srvURL(t, srv.URL)), RPM: 3, PeakRPM: 3, PeakInflight: 1}
|
||||
for range 2 {
|
||||
if s := g.Snapshot(); len(s) != 1 || s[0] != want {
|
||||
t.Fatalf("snapshot = %+v, want [%+v]", s, want)
|
||||
}
|
||||
}
|
||||
g.report()
|
||||
if s := g.Snapshot(); len(s) != 1 || s[0].RPM != 0 || s[0].PeakRPM != 3 {
|
||||
t.Fatalf("snapshot after report = %+v, want rpm 0 and peak 3", s)
|
||||
}
|
||||
}
|
||||
|
||||
// SetProfile never starts a goroutine of its own: one ticker serves every
|
||||
// governor.
|
||||
func TestSetProfileStartsNoGoroutine(t *testing.T) {
|
||||
keepProfile(t)
|
||||
SetProfile(Profiles["quiet"]) // the reporter is running from here on
|
||||
before := runtime.NumGoroutine()
|
||||
for range 50 {
|
||||
SetProfile(Profiles["balanced"])
|
||||
SetProfile(Profile{Name: "off"})
|
||||
}
|
||||
if after := runtime.NumGoroutine(); after > before {
|
||||
t.Fatalf("goroutines %d -> %d after 100 SetProfile calls", before, after)
|
||||
}
|
||||
}
|
||||
|
||||
func srvURL(t *testing.T, raw string) *url.URL {
|
||||
t.Helper()
|
||||
u, err := url.Parse(raw)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return u
|
||||
}
|
||||
+124
-22
@@ -1,6 +1,7 @@
|
||||
// Package egress is the governor the platform-owned HTTP requests to a target
|
||||
// pass through (WebFetch, the http custom tools, the enrich probe). A noise
|
||||
// profile (ARTEX_NOISE_PROFILE) sets every lever: a token bucket, a jitter and a
|
||||
// profile (the noise_profile setting, else ARTEX_NOISE_PROFILE; SetProfile
|
||||
// changes it at runtime) sets every lever: a token bucket, a jitter and a
|
||||
// concurrency cap per host, a global concurrency cap, backoff on 429/503, and,
|
||||
// for the GET requests of one task, coalescing and a short cache. The default
|
||||
// profile, off, is a strict passthrough: requests and timing are unchanged.
|
||||
@@ -22,6 +23,7 @@ import (
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
"unicode"
|
||||
|
||||
@@ -54,42 +56,91 @@ var Profiles = map[string]Profile{
|
||||
"aggressive": {Name: "aggressive", RPS: 20, Burst: 40, MaxPerHost: 4, MaxGlobal: 32, Jitter: 0.1, CacheTTL: 10 * time.Second, MaxRetries: 2, BaseDelay: 500 * time.Millisecond, MaxDelay: 10 * time.Second},
|
||||
}
|
||||
|
||||
// ProfileFromEnv returns the profile ARTEX_NOISE_PROFILE names; unset, "off" or
|
||||
// unknown is off.
|
||||
func ProfileFromEnv() Profile {
|
||||
name := strings.ToLower(strings.TrimSpace(os.Getenv("ARTEX_NOISE_PROFILE")))
|
||||
p, ok := Profiles[name]
|
||||
if !ok {
|
||||
if name != "" && name != "off" {
|
||||
log.Printf("[egress] ARTEX_NOISE_PROFILE=%q is not off, quiet, balanced or aggressive; using off", name)
|
||||
}
|
||||
return Profile{Name: "off"}
|
||||
// ProfileByName returns the profile name names (off, quiet, balanced or
|
||||
// aggressive, any case); ok is false for any other name, "" included. A profile
|
||||
// other than off carries ARTEX_TARGET_USER_AGENT; off stays a strict passthrough.
|
||||
func ProfileByName(name string) (p Profile, ok bool) {
|
||||
name = strings.ToLower(strings.TrimSpace(name))
|
||||
if name == "off" {
|
||||
return Profile{Name: "off"}, true
|
||||
}
|
||||
if p, ok = Profiles[name]; !ok {
|
||||
return Profile{Name: "off"}, false
|
||||
}
|
||||
if ua := strings.TrimSpace(os.Getenv("ARTEX_TARGET_USER_AGENT")); strings.ContainsFunc(ua, unicode.IsControl) {
|
||||
log.Printf("[egress] ARTEX_TARGET_USER_AGENT contains a control character; ignored")
|
||||
} else {
|
||||
p.UserAgent = ua
|
||||
}
|
||||
return p, true
|
||||
}
|
||||
|
||||
// ProfileFromEnv returns the profile ARTEX_NOISE_PROFILE names; unset, "off" or
|
||||
// unknown is off.
|
||||
func ProfileFromEnv() Profile {
|
||||
name := strings.TrimSpace(os.Getenv("ARTEX_NOISE_PROFILE"))
|
||||
p, ok := ProfileByName(name)
|
||||
if !ok && name != "" {
|
||||
log.Printf("[egress] ARTEX_NOISE_PROFILE=%q is not off, quiet, balanced or aggressive; using off", name)
|
||||
}
|
||||
return p
|
||||
}
|
||||
|
||||
// Shared is the one governor of the process, built from the environment on
|
||||
// first use; WebFetch, the http custom tools and the enrich probe all use it.
|
||||
// Under a profile it logs the noise counter of each host once a minute.
|
||||
var Shared = sync.OnceValue(func() *Governor {
|
||||
g := New(ProfileFromEnv())
|
||||
// current is the governor new requests use. SetProfile replaces it; a request
|
||||
// already in flight keeps the governor it started on, so its slots are
|
||||
// released where they were taken.
|
||||
var (
|
||||
current atomic.Pointer[Governor]
|
||||
initCurrent sync.Once
|
||||
reporterOnce sync.Once
|
||||
)
|
||||
|
||||
// Current returns the governor of the process: the one SetProfile published
|
||||
// last, or, before any SetProfile, one built from the environment on first use.
|
||||
// WebFetch, the http custom tools and the enrich probe resolve it per request.
|
||||
func Current() *Governor {
|
||||
if g := current.Load(); g != nil {
|
||||
return g
|
||||
}
|
||||
initCurrent.Do(func() {
|
||||
g := New(ProfileFromEnv())
|
||||
if current.CompareAndSwap(nil, g) {
|
||||
log.Printf("[egress] noise profile: %s", g.p.Name)
|
||||
}
|
||||
startReporter()
|
||||
})
|
||||
return current.Load()
|
||||
}
|
||||
|
||||
// CurrentProfile is the profile of Current.
|
||||
func CurrentProfile() Profile { return Current().p }
|
||||
|
||||
// SetProfile builds a governor for p and publishes it: new requests use it,
|
||||
// requests in flight finish on the previous one. The counters start afresh.
|
||||
func SetProfile(p Profile) *Governor {
|
||||
g := New(p)
|
||||
current.Store(g)
|
||||
log.Printf("[egress] noise profile: %s", g.p.Name)
|
||||
if !g.off() {
|
||||
startReporter()
|
||||
return g
|
||||
}
|
||||
|
||||
// startReporter starts the one package-level ticker that logs, once a minute,
|
||||
// the noise counter of each host of the current governor. SetProfile never
|
||||
// starts another.
|
||||
func startReporter() {
|
||||
reporterOnce.Do(func() {
|
||||
go func() {
|
||||
for range time.Tick(time.Minute) {
|
||||
for _, l := range g.report() {
|
||||
t := time.NewTicker(time.Minute)
|
||||
defer t.Stop()
|
||||
for range t.C {
|
||||
for _, l := range Current().report() {
|
||||
log.Print(l)
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
return g
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
// Governor holds the per-host state and the per-task response cache.
|
||||
type Governor struct {
|
||||
@@ -123,6 +174,9 @@ func New(p Profile) *Governor {
|
||||
|
||||
func (g *Governor) off() bool { return g.p.RPS <= 0 }
|
||||
|
||||
// Profile is the profile g was built for.
|
||||
func (g *Governor) Profile() Profile { return g.p }
|
||||
|
||||
// HostKey is the lowercased host, with the port when it is not the default.
|
||||
func HostKey(u *url.URL) string {
|
||||
h, p := strings.ToLower(u.Hostname()), u.Port()
|
||||
@@ -240,6 +294,35 @@ func (g *Governor) report() []string {
|
||||
return out
|
||||
}
|
||||
|
||||
// HostNoise is the noise counter of one host: the current minute (RPM,
|
||||
// DuplicatesAvoided, PeakInflight), the peak per minute so far, and the
|
||||
// requests in flight now.
|
||||
type HostNoise struct {
|
||||
Host string `json:"host"`
|
||||
RPM int `json:"rpm"`
|
||||
PeakRPM int `json:"peak_rpm"`
|
||||
PeakInflight int `json:"peak_inflight"`
|
||||
Inflight int `json:"inflight"`
|
||||
DuplicatesAvoided int `json:"duplicates_avoided"`
|
||||
}
|
||||
|
||||
// Snapshot returns the counter of every host that has sent or deduplicated a
|
||||
// request, sorted by host. Unlike report it starts no new minute.
|
||||
func (g *Governor) Snapshot() []HostNoise {
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
out := []HostNoise{}
|
||||
for key, h := range g.hosts {
|
||||
if h.sent == 0 && h.dupes == 0 && h.inflight == 0 && h.peakRPM == 0 {
|
||||
continue
|
||||
}
|
||||
out = append(out, HostNoise{Host: key, RPM: h.sent, PeakRPM: max(h.peakRPM, h.sent),
|
||||
PeakInflight: h.peakInflight, Inflight: h.inflight, DuplicatesAvoided: h.dupes})
|
||||
}
|
||||
sort.Slice(out, func(i, j int) bool { return out[i].Host < out[j].Host })
|
||||
return out
|
||||
}
|
||||
|
||||
var errNotShared = errors.New("egress: response not shared")
|
||||
|
||||
// share runs fn once for the concurrent callers of key and keeps its value for
|
||||
@@ -344,6 +427,25 @@ type transport struct {
|
||||
base http.RoundTripper
|
||||
}
|
||||
|
||||
// Dynamic returns base (nil: http.DefaultTransport) governed by the governor
|
||||
// current at each request, for a client built once and kept: a later
|
||||
// SetProfile applies to its next request. Off, the request goes to base as is.
|
||||
func Dynamic(namespace string, base http.RoundTripper) http.RoundTripper {
|
||||
if base == nil {
|
||||
base = http.DefaultTransport
|
||||
}
|
||||
return dynamicTransport{namespace, base}
|
||||
}
|
||||
|
||||
type dynamicTransport struct {
|
||||
namespace string
|
||||
base http.RoundTripper
|
||||
}
|
||||
|
||||
func (d dynamicTransport) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
return Current().Transport(d.namespace, d.base).RoundTrip(req)
|
||||
}
|
||||
|
||||
const maxSharedBody = 2 << 20 // a larger response is neither coalesced nor cached
|
||||
|
||||
func (t *transport) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
|
||||
@@ -0,0 +1,67 @@
|
||||
package egress
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"net/url"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Bash commands reach targets through the recording proxy, outside the clients
|
||||
// this package governs. With traffic capture on, the proxy URL in their
|
||||
// environment carries a marker userinfo (MarkProxyURL). The proxy paces each
|
||||
// request of a marked client connection with Wait, and leaves unmarked
|
||||
// connections alone: WebFetch, the http custom tools and the enrich probe use
|
||||
// the plain proxy URL and are already paced by their governor, so no request
|
||||
// is paced twice. Best effort, not a security boundary: a command that ignores
|
||||
// the proxy, or drops the marker, is not paced.
|
||||
|
||||
// PaceUser is the user and password of the marker userinfo.
|
||||
const PaceUser = "artex-pace"
|
||||
|
||||
// MarkProxyURL returns raw with the marker userinfo. raw is returned unchanged
|
||||
// when it does not parse or already carries a userinfo.
|
||||
func MarkProxyURL(raw string) string {
|
||||
u, err := url.Parse(raw)
|
||||
if err != nil || u.Host == "" || u.User != nil {
|
||||
return raw
|
||||
}
|
||||
u.User = url.UserPassword(PaceUser, PaceUser)
|
||||
return u.String()
|
||||
}
|
||||
|
||||
// IsPaceMarker reports whether a Proxy-Authorization value is Basic
|
||||
// credentials for PaceUser.
|
||||
func IsPaceMarker(proxyAuthorization string) bool {
|
||||
scheme, enc, ok := strings.Cut(strings.TrimSpace(proxyAuthorization), " ")
|
||||
if !ok || !strings.EqualFold(scheme, "Basic") {
|
||||
return false
|
||||
}
|
||||
dec, err := base64.StdEncoding.DecodeString(strings.TrimSpace(enc))
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
user, _, _ := strings.Cut(string(dec), ":")
|
||||
return user == PaceUser
|
||||
}
|
||||
|
||||
// Wait paces one request to hostKey under the current profile. See
|
||||
// Governor.Wait.
|
||||
func Wait(ctx context.Context, hostKey string) error { return Current().Wait(ctx, hostKey) }
|
||||
|
||||
// Wait takes a token of the host's bucket plus the jitter, as a governed request
|
||||
// does, and counts the request in the host's noise counter. It takes no
|
||||
// in-flight slot, and nothing is cached, coalesced or retried. Off, it returns
|
||||
// at once and counts nothing. An error means ctx ended first; the request is
|
||||
// counted either way, since the caller sends it regardless.
|
||||
func (g *Governor) Wait(ctx context.Context, hostKey string) error {
|
||||
if g.off() {
|
||||
return nil
|
||||
}
|
||||
h := g.host(hostKey)
|
||||
err := sleep(ctx, g.reserve(h)+g.jitter())
|
||||
g.mu.Lock()
|
||||
h.sent++
|
||||
g.mu.Unlock()
|
||||
return err
|
||||
}
|
||||
Executable
+85
@@ -0,0 +1,85 @@
|
||||
package egress
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestMarkProxyURL(t *testing.T) {
|
||||
for in, want := range map[string]string{
|
||||
"http://127.0.0.1:8788": "http://artex-pace:artex-pace@127.0.0.1:8788",
|
||||
"http://u:p@127.0.0.1:8788": "http://u:p@127.0.0.1:8788",
|
||||
"": "",
|
||||
"127.0.0.1:8788": "127.0.0.1:8788",
|
||||
"socks5://proxy.example.test:1080": "socks5://artex-pace:artex-pace@proxy.example.test:1080",
|
||||
} {
|
||||
if got := MarkProxyURL(in); got != want {
|
||||
t.Errorf("MarkProxyURL(%q) = %q, want %q", in, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsPaceMarker(t *testing.T) {
|
||||
enc := func(s string) string { return base64.StdEncoding.EncodeToString([]byte(s)) }
|
||||
for v, want := range map[string]bool{
|
||||
"Basic " + enc("artex-pace:artex-pace"): true,
|
||||
"basic " + enc("artex-pace:x"): true,
|
||||
"Basic " + enc("user:pass"): false,
|
||||
"Bearer " + enc("artex-pace:artex-pace"): false,
|
||||
"Basic !!!": false,
|
||||
"Basic": false,
|
||||
"": false,
|
||||
} {
|
||||
if got := IsPaceMarker(v); got != want {
|
||||
t.Errorf("IsPaceMarker(%q) = %v, want %v", v, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestWait(t *testing.T) {
|
||||
if err := New(Profile{Name: "off"}).Wait(t.Context(), "a.example"); err != nil {
|
||||
t.Fatalf("off Wait: %v", err)
|
||||
}
|
||||
off := New(Profile{Name: "off"})
|
||||
_ = off.Wait(t.Context(), "a.example")
|
||||
if s := off.Snapshot(); len(s) != 0 {
|
||||
t.Fatalf("off counted: %v", s)
|
||||
}
|
||||
|
||||
// 10 req/s, burst 1, no jitter: the second Wait takes ~100 ms.
|
||||
g := New(Profile{Name: "test", RPS: 10, Burst: 1, MaxPerHost: 1, MaxGlobal: 1})
|
||||
start := time.Now()
|
||||
for range 2 {
|
||||
if err := g.Wait(t.Context(), "a.example"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if d := time.Since(start); d < 80*time.Millisecond {
|
||||
t.Fatalf("two Waits at 10 req/s took %v", d)
|
||||
}
|
||||
if s := g.Snapshot(); len(s) != 1 || s[0].RPM != 2 || s[0].Inflight != 0 || s[0].PeakInflight != 0 {
|
||||
t.Fatalf("snapshot after Wait = %+v, want rpm 2 and no in-flight slot", s)
|
||||
}
|
||||
|
||||
// A wait past the context's end returns its error and still counts.
|
||||
slow := New(Profile{Name: "test", RPS: 0.5, Burst: 1, MaxPerHost: 1, MaxGlobal: 1})
|
||||
_ = slow.Wait(t.Context(), "b.example")
|
||||
ctx, cancel := context.WithTimeout(t.Context(), 20*time.Millisecond)
|
||||
defer cancel()
|
||||
if err := slow.Wait(ctx, "b.example"); !errors.Is(err, context.DeadlineExceeded) {
|
||||
t.Fatalf("bounded Wait: %v, want DeadlineExceeded", err)
|
||||
}
|
||||
if s := slow.Snapshot(); len(s) != 1 || s[0].RPM != 2 {
|
||||
t.Fatalf("snapshot after a timed-out Wait = %+v, want rpm 2", s)
|
||||
}
|
||||
|
||||
// The package Wait uses the current profile.
|
||||
keepProfile(t)
|
||||
SetProfile(Profile{Name: "off"})
|
||||
if err := Wait(t.Context(), "c.example"); err != nil || len(Current().Snapshot()) != 0 {
|
||||
t.Fatalf("package Wait under off: err=%v snapshot=%v", err, Current().Snapshot())
|
||||
}
|
||||
}
|
||||
+3
-2
@@ -99,7 +99,8 @@ func New(as *db.AssetStore, proxy func() string, workers int) *Engine {
|
||||
// probe without a scope is refused. When the scope's gate is on, the probe
|
||||
// connects to the target itself, without the proxy, so that the dialer decides
|
||||
// on the target and connects only to the addresses a pin covers. Every request
|
||||
// goes through the shared egress governor (the noise profile; off by default).
|
||||
// goes through the egress governor current at that request (the noise profile;
|
||||
// off by default), so a changed profile applies without rebuilding the client.
|
||||
func buildClient(proxy func() string) *http.Client {
|
||||
tr := &http.Transport{
|
||||
TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
|
||||
@@ -120,7 +121,7 @@ func buildClient(proxy func() string) *http.Client {
|
||||
}
|
||||
}
|
||||
return &http.Client{
|
||||
Transport: egress.Shared().Transport("enrich", tr),
|
||||
Transport: egress.Dynamic("enrich", tr),
|
||||
Timeout: httpTimeout,
|
||||
// cap redirects; keep them within scope by re-checking at probe time
|
||||
CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse },
|
||||
|
||||
+6
-6
@@ -195,11 +195,11 @@ func (r *Recorder) Stream(ctx context.Context, req llm.CompletionRequest) iter.S
|
||||
return
|
||||
}
|
||||
// Always track usage (cheap); accumulate text/thinking only for bodies.
|
||||
// Anthropic (and the other providers) split token usage across events:
|
||||
// message_start carries ONLY the input side (input + cache), message_delta
|
||||
// ONLY the output. They must be FOLDED with Add — overwriting on delta
|
||||
// would zero out the input already counted at start (mirrors the SDK's own
|
||||
// llm.Accumulator; see norma/llm/accumulate.go).
|
||||
// message_start carries the input side (input + cache) and is added.
|
||||
// Since Norma v0.4.7 message_delta also echoes that input side next to the
|
||||
// output, so it is folded with llm.Usage.MergeEcho (max on echoed fields,
|
||||
// sum on output), as llm.Accumulator does; Add would count the input and
|
||||
// cache twice. An output-only delta folds the same either way.
|
||||
switch ev.Type {
|
||||
case llm.SETextDelta:
|
||||
if recordBodies {
|
||||
@@ -215,7 +215,7 @@ func (r *Recorder) Stream(ctx context.Context, req llm.CompletionRequest) iter.S
|
||||
if ev.StopReason != "" {
|
||||
stopReason = ev.StopReason
|
||||
}
|
||||
usage.Add(ev.Usage)
|
||||
usage.MergeEcho(ev.Usage)
|
||||
}
|
||||
if !yield(ev, err) {
|
||||
return
|
||||
|
||||
@@ -0,0 +1,118 @@
|
||||
package llmrec
|
||||
|
||||
import (
|
||||
"context"
|
||||
"iter"
|
||||
"testing"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// echoStreamProvider streams one message_start and one message_delta with the
|
||||
// given (already normalized) usage, as Norma's providers emit them.
|
||||
type echoStreamProvider struct {
|
||||
completeProvider
|
||||
start, delta llm.Usage
|
||||
}
|
||||
|
||||
func (p echoStreamProvider) Stream(context.Context, llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
|
||||
return func(yield func(llm.StreamEvent, error) bool) {
|
||||
for _, ev := range []llm.StreamEvent{
|
||||
{Type: llm.SEMessageStart, Usage: p.start},
|
||||
{Type: llm.SETextDelta, Text: "answer"},
|
||||
{Type: llm.SEMessageDelta, StopReason: "end_turn", Usage: p.delta},
|
||||
{Type: llm.SEMessageStop},
|
||||
} {
|
||||
if !yield(ev, nil) {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Norma normalizes Anthropic usage so InputTokens includes the cache read and
|
||||
// write subsets, and since v0.4.7 message_delta echoes that input side.
|
||||
var usageEchoCases = []struct {
|
||||
name string
|
||||
start, delta llm.Usage
|
||||
want llm.Usage
|
||||
}{
|
||||
{
|
||||
name: "full_echo",
|
||||
start: llm.Usage{InputTokens: 23, CacheReadTokens: 5, CacheWriteTokens: 3},
|
||||
delta: llm.Usage{InputTokens: 23, CacheReadTokens: 5, CacheWriteTokens: 3, OutputTokens: 5},
|
||||
want: llm.Usage{InputTokens: 23, OutputTokens: 5, CacheReadTokens: 5, CacheWriteTokens: 3},
|
||||
},
|
||||
{
|
||||
// Raw start: input 500, cache read 200, cache creation 40 → normalized 740.
|
||||
// Raw delta: input 500, output 77, no cache fields → normalized 500.
|
||||
name: "partial_echo",
|
||||
start: llm.Usage{InputTokens: 740, CacheReadTokens: 200, CacheWriteTokens: 40},
|
||||
delta: llm.Usage{InputTokens: 500, OutputTokens: 77},
|
||||
want: llm.Usage{InputTokens: 740, OutputTokens: 77, CacheReadTokens: 200, CacheWriteTokens: 40},
|
||||
},
|
||||
{
|
||||
name: "no_echo",
|
||||
start: llm.Usage{InputTokens: 23, CacheReadTokens: 5, CacheWriteTokens: 3},
|
||||
delta: llm.Usage{OutputTokens: 5},
|
||||
want: llm.Usage{InputTokens: 23, OutputTokens: 5, CacheReadTokens: 5, CacheWriteTokens: 3},
|
||||
},
|
||||
}
|
||||
|
||||
func drainRecorderStream(t *testing.T, ctx context.Context, r *Recorder) {
|
||||
t.Helper()
|
||||
for _, err := range r.Stream(ctx, llm.CompletionRequest{}) {
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The metered usage of a streamed call counts an echoed input side once.
|
||||
func TestStreamUsageEchoCountedOnce(t *testing.T) {
|
||||
for _, tc := range usageEchoCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
var got []llm.Usage
|
||||
ctx := WithUsageSink(t.Context(), func(u llm.Usage) { got = append(got, u) })
|
||||
drainRecorderStream(t, ctx, Wrap(echoStreamProvider{start: tc.start, delta: tc.delta}, nil, "fixture", "profile", "", "", func() bool { return false }))
|
||||
if len(got) != 1 || got[0] != tc.want {
|
||||
t.Fatalf("metered %+v, want one call with %+v", got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// The llm_usage ledger gets one row per streamed call, with the echo counted once.
|
||||
func TestStreamUsageEchoLedgerRow(t *testing.T) {
|
||||
dsn, _, err := db.DSN()
|
||||
if err != nil {
|
||||
t.Skipf("no database config (%v)", err)
|
||||
}
|
||||
pg, err := db.Open(dsn)
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v)", err)
|
||||
}
|
||||
defer pg.Close()
|
||||
if err := pg.EnsureLLMUsageTable(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, tc := range usageEchoCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
profile := "usage-echo-" + uuid.NewString()
|
||||
t.Cleanup(func() { _, _ = pg.Exec(`DELETE FROM llm_usage WHERE profile_name=$1`, profile) })
|
||||
drainRecorderStream(t, t.Context(), Wrap(echoStreamProvider{start: tc.start, delta: tc.delta}, pg, "fixture", profile, "", "", func() bool { return false }))
|
||||
var rows int
|
||||
var got llm.Usage
|
||||
err := pg.QueryRow(`SELECT count(*), COALESCE(sum(input_tokens),0), COALESCE(sum(output_tokens),0), COALESCE(sum(cache_read),0), COALESCE(sum(cache_write),0)
|
||||
FROM llm_usage WHERE profile_name=$1`, profile).Scan(&rows, &got.InputTokens, &got.OutputTokens, &got.CacheReadTokens, &got.CacheWriteTokens)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if rows != 1 || got != tc.want {
|
||||
t.Fatalf("ledger rows=%d usage=%+v, want 1 row with %+v", rows, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -6,9 +6,9 @@ import (
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
)
|
||||
|
||||
// hasSkillTool reports whether the packed tool set contains the Skill meta-tool.
|
||||
|
||||
@@ -566,13 +566,14 @@ func (s *Server) runHTTPTool(ctx context.Context, key string, execRaw json.RawMe
|
||||
client := &http.Client{Timeout: timeoutOr(spec.TimeoutMs, 30000)}
|
||||
// Every redirect hop goes through the target gate, as the rendered URL did.
|
||||
client.CheckRedirect = intercept.RedirectCheck(key)
|
||||
// The shared egress governor paces every request and redirect hop per host
|
||||
// under the noise profile (off by default: the base transport itself).
|
||||
// The current egress governor paces every request and redirect hop per host
|
||||
// under the noise profile (off by default: the base transport itself). It is
|
||||
// resolved per request, so a changed profile applies to the next call.
|
||||
var base http.RoundTripper = http.DefaultTransport
|
||||
if tr := s.httpToolTransport(ctx, key, spec); tr != nil {
|
||||
base = tr
|
||||
}
|
||||
client.Transport = egress.Shared().Transport(key, base)
|
||||
client.Transport = egress.Current().Transport(key, base)
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return actool.Errorf("request failed: " + err.Error()), nil
|
||||
|
||||
@@ -7,19 +7,19 @@ import (
|
||||
|
||||
func TestGoalGateCheck(t *testing.T) {
|
||||
good := map[string]string{
|
||||
"url target": "Audit " + strings.Repeat("x ", 60) + " at https://example.com/app every module for injection issues and record one finding per issue with evidence.",
|
||||
"bare domain": strings.Repeat("probe ", 25) + "the api at target.example.org for unauthenticated data exposure, one record proof, then stop.",
|
||||
"ipv4": strings.Repeat("recon ", 25) + "10.0.0.5 and map every listening service before any test.",
|
||||
"localhost port": strings.Repeat("verify ", 25) + "localhost:3000 health endpoints respond identically before and after config reload.",
|
||||
"path": strings.Repeat("review ", 25) + "the module at /app/data/core for hardcoded secrets and report each hit.",
|
||||
"targetless opt-out": "[targetless] " + strings.Repeat("summarize ", 20) + "the attached material into a prioritized hypothesis list.",
|
||||
"url target": "Audit " + strings.Repeat("x ", 60) + " at https://example.com/app every module for injection issues and record one finding per issue with evidence.",
|
||||
"bare domain": strings.Repeat("probe ", 25) + "the api at target.example.org for unauthenticated data exposure, one record proof, then stop.",
|
||||
"ipv4": strings.Repeat("recon ", 25) + "10.0.0.5 and map every listening service before any test.",
|
||||
"localhost port": strings.Repeat("verify ", 25) + "localhost:3000 health endpoints respond identically before and after config reload.",
|
||||
"path": strings.Repeat("review ", 25) + "the module at /app/data/core for hardcoded secrets and report each hit.",
|
||||
"targetless opt-out": "[targetless] " + strings.Repeat("summarize ", 20) + "the attached material into a prioritized hypothesis list.",
|
||||
}
|
||||
good["targetless short but marked"] = "[targetless] analyze nothing"
|
||||
|
||||
bad := map[string]string{
|
||||
"no target long": strings.Repeat("do a thorough security review of everything and report anything interesting that comes up during the exploration phase of this audit, thanks. ", 2),
|
||||
"too short": "audit the site",
|
||||
"short and bare": "find bugs",
|
||||
"no target long": strings.Repeat("do a thorough security review of everything and report anything interesting that comes up during the exploration phase of this audit, thanks. ", 2),
|
||||
"too short": "audit the site",
|
||||
"short and bare": "find bugs",
|
||||
}
|
||||
for name, g := range good {
|
||||
if name == "targetless short+:url" {
|
||||
|
||||
+12
-7
@@ -19,6 +19,7 @@ import (
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/config"
|
||||
pgdb "github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/egress"
|
||||
"github.com/Autumn-27/artex/enrich"
|
||||
"github.com/Autumn-27/artex/guard"
|
||||
"github.com/Autumn-27/artex/intercept"
|
||||
@@ -54,13 +55,13 @@ type Task struct {
|
||||
LLMChainRevision int64 `json:"-"`
|
||||
LLMFailoverState string `json:"llm_failover_state,omitempty"`
|
||||
// LLMRotate (#85): per-intent round-robin over the task's LLM chain heads.
|
||||
LLMRotate bool `json:"llm_rotate,omitempty"`
|
||||
rotateMu sync.Mutex
|
||||
rotateNext int
|
||||
LLMFailoverReason string `json:"llm_failover_reason,omitempty"`
|
||||
SourceTaskIDs []int64 `json:"source_task_ids,omitempty"`
|
||||
CompanyIDs []int64 `json:"company_ids,omitempty"`
|
||||
Status string `json:"status"` // persisted lifecycle status (done/failed/timeout are terminal; empty/other is derived from runtime state)
|
||||
LLMRotate bool `json:"llm_rotate,omitempty"`
|
||||
rotateMu sync.Mutex
|
||||
rotateNext int
|
||||
LLMFailoverReason string `json:"llm_failover_reason,omitempty"`
|
||||
SourceTaskIDs []int64 `json:"source_task_ids,omitempty"`
|
||||
CompanyIDs []int64 `json:"company_ids,omitempty"`
|
||||
Status string `json:"status"` // persisted lifecycle status (done/failed/timeout are terminal; empty/other is derived from runtime state)
|
||||
// Task-level timeout (see docs/task-timeout-and-settle-design.md). DeadlineAt/FirstRunAt are unix seconds, 0 = unset/not yet run.
|
||||
TimeoutSeconds int `json:"timeout_seconds"`
|
||||
PlanHeartbeatSeconds int `json:"plan_heartbeat_seconds"` // planner heartbeat trigger interval (seconds)
|
||||
@@ -304,6 +305,9 @@ const (
|
||||
// when on, the four agent kinds wired into the platform (planner/worker/main
|
||||
// agent/chat) have context compaction taken over by noa, replacing the built-in compaction. Read once per run; toggling affects only runs started afterwards.
|
||||
settingNoaCompaction = "noa_compaction"
|
||||
// Target noise profile (off|quiet|balanced|aggressive). When stored it wins over
|
||||
// ARTEX_NOISE_PROFILE, which stays the fallback; applied at startup and on PUT.
|
||||
settingNoiseProfile = "noise_profile"
|
||||
// defaultWebSearchBackend is used when web search is on but no backend was picked.
|
||||
defaultWebSearchBackend = "ddgs"
|
||||
// deepSeekWebSearchBackend borrows the active LLM profile instead of its own
|
||||
@@ -401,6 +405,7 @@ func NewManager(dir, proxyAddr string) (*Manager, error) {
|
||||
pg.Close()
|
||||
return nil, fmt.Errorf("recover finding retests: %w", err)
|
||||
}
|
||||
egress.SetProfile(startupNoiseProfile(pg))
|
||||
if err := pg.EnsureLLMRecordsTable(); err != nil {
|
||||
log.Printf("[llmrec] create table: %v", err)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,32 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"log"
|
||||
"net/http"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/egress"
|
||||
)
|
||||
|
||||
// startupNoiseProfile is the noise profile to apply at startup: the stored
|
||||
// noise_profile setting when there is a valid one, else ARTEX_NOISE_PROFILE.
|
||||
func startupNoiseProfile(pg *db.DB) egress.Profile {
|
||||
v, ok, err := pg.GetSetting(settingNoiseProfile)
|
||||
if err != nil {
|
||||
log.Printf("[egress] read %s setting: %v; using ARTEX_NOISE_PROFILE", settingNoiseProfile, err)
|
||||
}
|
||||
if ok {
|
||||
if p, valid := egress.ProfileByName(v); valid {
|
||||
return p
|
||||
}
|
||||
log.Printf("[egress] stored %s=%q is not off, quiet, balanced or aggressive; using ARTEX_NOISE_PROFILE", settingNoiseProfile, v)
|
||||
}
|
||||
return egress.ProfileFromEnv()
|
||||
}
|
||||
|
||||
// noiseStatus returns the current noise profile and the per-host counters of
|
||||
// its governor, read without starting a new minute.
|
||||
func (s *Server) noiseStatus(w http.ResponseWriter, r *http.Request) {
|
||||
g := egress.Current()
|
||||
writeJSON(w, 200, map[string]any{"profile": g.Profile().Name, "hosts": g.Snapshot()})
|
||||
}
|
||||
@@ -0,0 +1,125 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Autumn-27/artex/egress"
|
||||
)
|
||||
|
||||
func TestNoiseProfileSettingHTTP(t *testing.T) {
|
||||
skipWithoutDSN(t)
|
||||
prev := egress.CurrentProfile()
|
||||
t.Cleanup(func() { egress.SetProfile(prev) })
|
||||
t.Setenv("ARTEX_NOISE_PROFILE", "quiet")
|
||||
t.Setenv("ARTEX_TARGET_USER_AGENT", "artex-noise-test/1")
|
||||
|
||||
m, err := NewManager(t.TempDir(), "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer m.Close()
|
||||
clear := func() {
|
||||
if _, err := m.pg.Exec(`DELETE FROM settings WHERE key=$1`, settingNoiseProfile); err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
}
|
||||
clear()
|
||||
defer clear() // before m.Close
|
||||
|
||||
// Precedence: no setting -> env; a valid setting wins; an invalid one falls back to env.
|
||||
if p := startupNoiseProfile(m.pg); p.Name != "quiet" {
|
||||
t.Fatalf("no setting: %q, want env quiet", p.Name)
|
||||
}
|
||||
if err := m.pg.SetSetting(settingNoiseProfile, "balanced"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
m2, err := NewManager(t.TempDir(), "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
m2.Close()
|
||||
if p := egress.CurrentProfile(); p.Name != "balanced" || p.UserAgent != "artex-noise-test/1" {
|
||||
t.Fatalf("NewManager with setting balanced applied %+v", p)
|
||||
}
|
||||
if err := m.pg.SetSetting(settingNoiseProfile, "loud"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if p := startupNoiseProfile(m.pg); p.Name != "quiet" {
|
||||
t.Fatalf("invalid setting: %q, want env quiet", p.Name)
|
||||
}
|
||||
if err := m.pg.SetSetting(settingNoiseProfile, "balanced"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
s := New(ctx, m, t.TempDir(), t.TempDir(), t.TempDir())
|
||||
defer stopTestServer(t, s, cancel)
|
||||
token, err := signJWT(s.jwtKey)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
do := func(method, path, body string) (int, map[string]any) {
|
||||
t.Helper()
|
||||
req := httptest.NewRequest(method, path, strings.NewReader(body))
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
rec := httptest.NewRecorder()
|
||||
s.Handler().ServeHTTP(rec, req)
|
||||
var out map[string]any
|
||||
_ = json.Unmarshal(rec.Body.Bytes(), &out)
|
||||
return rec.Code, out
|
||||
}
|
||||
egress.SetProfile(egress.Profiles["balanced"])
|
||||
if code, out := do(http.MethodGet, "/api/settings", ""); code != 200 || out["noise_profile"] != "balanced" {
|
||||
t.Fatalf("GET settings: %d noise_profile=%v", code, out["noise_profile"])
|
||||
}
|
||||
if code, out := do(http.MethodPut, "/api/settings", `{"noise_profile":"Aggressive"}`); code != 200 || out["noise_profile"] != "aggressive" {
|
||||
t.Fatalf("PUT aggressive: %d noise_profile=%v", code, out["noise_profile"])
|
||||
}
|
||||
if v, _, _ := m.pg.GetSetting(settingNoiseProfile); v != "aggressive" || egress.CurrentProfile().Name != "aggressive" {
|
||||
t.Fatalf("after PUT: stored %q, current %q", v, egress.CurrentProfile().Name)
|
||||
}
|
||||
for _, bad := range []string{`{"noise_profile":"loud"}`, `{"noise_profile":""}`, `{"noise_profile":1}`} {
|
||||
if code, _ := do(http.MethodPut, "/api/settings", bad); code != 400 {
|
||||
t.Errorf("PUT %s: %d, want 400", bad, code)
|
||||
}
|
||||
}
|
||||
if code, out := do(http.MethodPut, "/api/settings", `{"noise_profile":null}`); code != 200 || out["noise_profile"] != "aggressive" {
|
||||
t.Fatalf("PUT null: %d noise_profile=%v, want unchanged", code, out["noise_profile"])
|
||||
}
|
||||
if v, _, _ := m.pg.GetSetting(settingNoiseProfile); v != "aggressive" || egress.CurrentProfile().Name != "aggressive" {
|
||||
t.Fatalf("after invalid and null PUTs: stored %q, current %q", v, egress.CurrentProfile().Name)
|
||||
}
|
||||
|
||||
// The noise endpoint reads the current governor's counters.
|
||||
target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
|
||||
defer target.Close()
|
||||
req, _ := http.NewRequest(http.MethodGet, target.URL, nil)
|
||||
resp, err := egress.Current().Transport("noise-test", nil).RoundTrip(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
resp.Body.Close()
|
||||
code, out := do(http.MethodGet, "/api/settings/noise", "")
|
||||
hosts, _ := out["hosts"].([]any)
|
||||
if code != 200 || out["profile"] != "aggressive" || len(hosts) != 1 {
|
||||
t.Fatalf("GET noise: %d %v", code, out)
|
||||
}
|
||||
host, _ := hosts[0].(map[string]any)
|
||||
for _, key := range []string{"host", "rpm", "peak_rpm", "peak_inflight", "inflight", "duplicates_avoided"} {
|
||||
if _, ok := host[key]; !ok {
|
||||
t.Errorf("noise host misses %q: %v", key, host)
|
||||
}
|
||||
}
|
||||
if host["host"] != strings.TrimPrefix(target.URL, "http://") || host["rpm"] != float64(1) || host["inflight"] != float64(0) {
|
||||
t.Fatalf("noise host = %v", host)
|
||||
}
|
||||
// A second read does not reset the minute.
|
||||
if _, again := do(http.MethodGet, "/api/settings/noise", ""); again["hosts"].([]any)[0].(map[string]any)["rpm"] != float64(1) {
|
||||
t.Fatalf("second GET noise reset the counter: %v", again)
|
||||
}
|
||||
}
|
||||
+27
-7
@@ -22,6 +22,7 @@ import (
|
||||
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/egress"
|
||||
"github.com/Autumn-27/artex/intercept"
|
||||
"github.com/Autumn-27/artex/llmpool"
|
||||
"github.com/Autumn-27/artex/llmrec"
|
||||
@@ -833,6 +834,7 @@ func (s *Server) Handler() http.Handler {
|
||||
mux.HandleFunc("GET /api/llm/records/{id}", s.pgGetLLMRecord)
|
||||
mux.HandleFunc("GET /api/settings", s.getSettings)
|
||||
mux.HandleFunc("PUT /api/settings", s.putSettings)
|
||||
mux.HandleFunc("GET /api/settings/noise", s.noiseStatus) // current noise profile + per-host counters (no reset)
|
||||
// Vulnerability IM push. Channels are "multi-instance + per-instance filter rules" resources, so they get their own group of
|
||||
// REST endpoints instead of being stuffed into the flat /api/settings key-value store.
|
||||
mux.HandleFunc("GET /api/notify/meta", s.notifyMeta)
|
||||
@@ -1607,12 +1609,12 @@ type createTaskReq struct {
|
||||
CategoryID *int64 `json:"category_id,omitempty"`
|
||||
Description string `json:"description"`
|
||||
Goal string `json:"goal"`
|
||||
LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // the LLM profile to run this task on; omitted/null = use the active profile
|
||||
LLMProfileIDs []int64 `json:"llm_profile_ids,omitempty"` // ordered task-level profile chain; the first entry takes effect initially
|
||||
SourceTaskIDs []string `json:"source_task_ids,omitempty"` // direct, read-only inherited source tasks
|
||||
CompanyIDs []int64 `json:"company_ids,omitempty"` // links the company scope and snapshots its current assets; does not copy assets or force intent generation
|
||||
LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // the LLM profile to run this task on; omitted/null = use the active profile
|
||||
LLMProfileIDs []int64 `json:"llm_profile_ids,omitempty"` // ordered task-level profile chain; the first entry takes effect initially
|
||||
SourceTaskIDs []string `json:"source_task_ids,omitempty"` // direct, read-only inherited source tasks
|
||||
CompanyIDs []int64 `json:"company_ids,omitempty"` // links the company scope and snapshots its current assets; does not copy assets or force intent generation
|
||||
TimeoutSeconds int `json:"timeout_seconds"`
|
||||
LLMRotate *bool `json:"llm_rotate"` // task-level timeout (seconds); 0/omitted = no limit
|
||||
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
|
||||
@@ -1925,11 +1927,11 @@ func (s *Server) getTask(w http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
dto := taskDTO(t, s.resolvedTaskStatus(t))
|
||||
metrics, err := s.m.PG().TaskListMetricsAll()
|
||||
metrics, err := s.m.PG().TaskListMetrics(t.ExpID)
|
||||
if err != nil {
|
||||
log.Printf("[tasks] detail metrics query failed for %s — metrics zeroed: %v", t.ID, err)
|
||||
}
|
||||
s.applyTaskMetrics(&dto, t, metrics[t.ExpID])
|
||||
s.applyTaskMetrics(&dto, t, metrics)
|
||||
archiveBlockers, _ := s.m.PG().TaskArchiveBlockers()
|
||||
applyTaskArchiveBlocker(&dto, archiveBlockers)
|
||||
writeJSON(w, 200, dto)
|
||||
@@ -3528,6 +3530,8 @@ func (s *Server) settingsPayload() map[string]any {
|
||||
// Render limits of IM messages; 0 = no truncation.
|
||||
"notify_summary_max_runes": limits.SummaryRunes,
|
||||
"notify_assets_max_shown": limits.AssetsShown,
|
||||
// Effective target noise profile (off by default): the stored setting, else ARTEX_NOISE_PROFILE.
|
||||
"noise_profile": egress.CurrentProfile().Name,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3603,11 +3607,27 @@ func (s *Server) putSettings(w http.ResponseWriter, r *http.Request) {
|
||||
// Render limits of IM messages: summary characters (0-4000) and asset names listed (0-50); 0 = no truncation.
|
||||
NotifySummaryMaxRunes *int `json:"notify_summary_max_runes"`
|
||||
NotifyAssetsMaxShown *int `json:"notify_assets_max_shown"`
|
||||
// Target noise profile: off|quiet|balanced|aggressive; null = unchanged. Applies to the next request, no agent rebuild.
|
||||
NoiseProfile *string `json:"noise_profile"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
// Validated before any other field is written, so a bad value changes nothing.
|
||||
var noise egress.Profile
|
||||
if req.NoiseProfile != nil {
|
||||
var ok bool
|
||||
if noise, ok = egress.ProfileByName(*req.NoiseProfile); !ok {
|
||||
writeErr(w, 400, "noise_profile must be off, quiet, balanced or aggressive")
|
||||
return
|
||||
}
|
||||
if err := s.m.pg.SetSetting(settingNoiseProfile, noise.Name); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
egress.SetProfile(noise)
|
||||
}
|
||||
// A masked placeholder echoed back keeps the stored secret, like null.
|
||||
req.BraveKey, req.TavilyKey = keepIfMasked(req.BraveKey), keepIfMasked(req.TavilyKey)
|
||||
req.WebSearchProxy, req.GlobalProxy = keepIfMasked(req.WebSearchProxy), keepIfMasked(req.GlobalProxy)
|
||||
|
||||
@@ -38,6 +38,21 @@ VALUES ($1,'goal','{}','met'), ($1,'goal','{}','open')`, task.ExpID); err != nil
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// A second task, so the detail must pick its own row and the list must keep
|
||||
// serving every task.
|
||||
other, err := m.CreateTask("detail metrics other", "goal", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
otherID, _ := strconv.ParseInt(other.ID, 10, 64)
|
||||
defer func() { _, _ = m.pg.DeleteTaskCascade(otherID, false, true) }()
|
||||
defer func() { _, _ = m.pg.Exec(`DELETE FROM findings WHERE task_id=$1`, taskID) }()
|
||||
if _, err := m.pg.Exec(`
|
||||
INSERT INTO findings(task_id, severity) VALUES ($1,'critical'), ($1,'high'), ($1,'high'), ($1,'low'), ($2,'medium')`,
|
||||
taskID, otherID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
s := New(ctx, m, t.TempDir(), t.TempDir(), t.TempDir())
|
||||
defer stopTestServer(t, s, cancel)
|
||||
@@ -45,7 +60,19 @@ VALUES ($1,'goal','{}','met'), ($1,'goal','{}','open')`, task.ExpID); err != nil
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// After New: restoring the task resets running intents to open.
|
||||
if _, err := m.pg.Exec(`
|
||||
INSERT INTO exploration_nodes(exploration_id, kind, payload, state)
|
||||
VALUES ($1,'intent','{}','running'), ($1,'intent','{}','done')`, task.ExpID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
type findingCounts struct {
|
||||
Critical int `json:"critical"`
|
||||
High int `json:"high"`
|
||||
Medium int `json:"medium"`
|
||||
Low int `json:"low"`
|
||||
}
|
||||
var probe struct {
|
||||
Tokens struct {
|
||||
InputTokens int `json:"input_tokens"`
|
||||
@@ -53,8 +80,10 @@ VALUES ($1,'goal','{}','met'), ($1,'goal','{}','open')`, task.ExpID); err != nil
|
||||
CacheReadTokens int `json:"cache_read_tokens"`
|
||||
CacheWriteTokens int `json:"cache_write_tokens"`
|
||||
} `json:"tokens"`
|
||||
GoalsMet int `json:"goals_met"`
|
||||
GoalsTotal int `json:"goals_total"`
|
||||
GoalsMet int `json:"goals_met"`
|
||||
GoalsTotal int `json:"goals_total"`
|
||||
InFlight int `json:"in_flight"`
|
||||
Findings findingCounts `json:"findings"`
|
||||
}
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/tasks/"+task.ID, nil)
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
@@ -73,6 +102,9 @@ VALUES ($1,'goal','{}','met'), ($1,'goal','{}','open')`, task.ExpID); err != nil
|
||||
if probe.GoalsMet != 1 || probe.GoalsTotal != 2 {
|
||||
t.Fatalf("detail goal progress not filled from metrics: %d/%d", probe.GoalsMet, probe.GoalsTotal)
|
||||
}
|
||||
if probe.InFlight != 1 || probe.Findings != (findingCounts{Critical: 1, High: 2, Low: 1}) {
|
||||
t.Fatalf("detail in-flight/findings not filled from this task's metrics: %d %+v", probe.InFlight, probe.Findings)
|
||||
}
|
||||
|
||||
// The list endpoint must report the same numbers for the same task.
|
||||
listReq := httptest.NewRequest(http.MethodGet, "/api/tasks", nil)
|
||||
@@ -89,27 +121,32 @@ VALUES ($1,'goal','{}','met'), ($1,'goal','{}','open')`, task.ExpID); err != nil
|
||||
InputTokens int `json:"input_tokens"`
|
||||
OutputTokens int `json:"output_tokens"`
|
||||
} `json:"tokens"`
|
||||
GoalsMet int `json:"goals_met"`
|
||||
GoalsTotal int `json:"goals_total"`
|
||||
GoalsMet int `json:"goals_met"`
|
||||
GoalsTotal int `json:"goals_total"`
|
||||
InFlight int `json:"in_flight"`
|
||||
Findings findingCounts `json:"findings"`
|
||||
} `json:"tasks"`
|
||||
}
|
||||
if err := json.Unmarshal(listRec.Body.Bytes(), &listed); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
found := false
|
||||
found, foundOther := false, false
|
||||
for _, lt := range listed.Tasks {
|
||||
if lt.ID == other.ID {
|
||||
foundOther = lt.Findings == (findingCounts{Medium: 1})
|
||||
}
|
||||
if lt.ID != task.ID {
|
||||
continue
|
||||
}
|
||||
found = true
|
||||
if lt.Tokens.InputTokens != probe.Tokens.InputTokens ||
|
||||
lt.Tokens.OutputTokens != probe.Tokens.OutputTokens ||
|
||||
lt.GoalsMet != probe.GoalsMet || lt.GoalsTotal != probe.GoalsTotal {
|
||||
t.Fatalf("list and detail metrics disagree: list=%+v/%d detail=%+v/%d",
|
||||
lt.Tokens, lt.GoalsMet, probe.Tokens, probe.GoalsTotal)
|
||||
lt.GoalsMet != probe.GoalsMet || lt.GoalsTotal != probe.GoalsTotal ||
|
||||
lt.InFlight != probe.InFlight || lt.Findings != probe.Findings {
|
||||
t.Fatalf("list and detail metrics disagree: list=%+v detail=%+v", lt, probe)
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatalf("task %s missing from list response", task.ID)
|
||||
if !found || !foundOther {
|
||||
t.Fatalf("list misses a task or its metrics: task %s found=%v, task %s with its findings=%v", task.ID, found, other.ID, foundOther)
|
||||
}
|
||||
}
|
||||
|
||||
+75
-19
@@ -317,10 +317,14 @@ func (r *taskLLMRuntime) Complete(ctx context.Context, req llm.CompletionRequest
|
||||
}
|
||||
|
||||
func completeTaskLLM(ctx context.Context, taskID string, req llm.CompletionRequest, hooks taskLLMStreamHooks) (llm.Message, string, llm.Usage, error) {
|
||||
// total is the usage of the logical call: the sum of every physical attempt
|
||||
// (same-provider retries and profile failover). Each attempt is already billed
|
||||
// once by the recorder below; this sum only feeds the caller's accounting.
|
||||
var total llm.Usage
|
||||
for {
|
||||
selection, err := hooks.current(ctx)
|
||||
if err != nil {
|
||||
return llm.Message{}, "", llm.Usage{}, err
|
||||
return llm.Message{}, "", total, err
|
||||
}
|
||||
var (
|
||||
msg llm.Message
|
||||
@@ -333,13 +337,14 @@ func completeTaskLLM(ctx context.Context, taskID string, req llm.CompletionReque
|
||||
retries, backoffOf := sameProviderRetryPolicy(selection.retry)
|
||||
for attempt := 0; ; attempt++ {
|
||||
msg, sr, usage, callErr = selection.provider.Complete(ctx, req)
|
||||
total.Add(usage)
|
||||
if callErr != nil && ctx.Err() == nil &&
|
||||
attempt < retries && isRetryableStreamError(callErr) {
|
||||
backoff := backoffOf(attempt)
|
||||
log.Printf("[task-llm] task %s non-streaming call failed, retrying on the same provider after %v (%d/%d): %v",
|
||||
taskID, backoff, attempt+1, retries, callErr)
|
||||
if sleepCtx(ctx, backoff) {
|
||||
break // ctx cancelled during backoff → stop retrying
|
||||
return llm.Message{}, "", total, ctx.Err() // cancelled during backoff → stop
|
||||
}
|
||||
continue
|
||||
}
|
||||
@@ -347,7 +352,7 @@ func completeTaskLLM(ctx context.Context, taskID string, req llm.CompletionReque
|
||||
}
|
||||
if callErr == nil {
|
||||
resetLLMFailStreak(taskID, selection.profileID)
|
||||
return msg, sr, usage, nil
|
||||
return msg, sr, total, nil
|
||||
}
|
||||
// Non-quota failures on a chain profile escalate after a consecutive
|
||||
// streak (#99): same advance path as quota exhaust, truthful reason.
|
||||
@@ -356,31 +361,31 @@ func completeTaskLLM(ctx context.Context, taskID string, req llm.CompletionReque
|
||||
cause := fmt.Errorf("profile #%d failed %d consecutive calls (last: %w)", selection.profileID, llmFailStreakThreshold, callErr)
|
||||
transition, markErr := hooks.exhaust(selection, cause)
|
||||
if markErr != nil {
|
||||
return llm.Message{}, "", llm.Usage{}, fmt.Errorf("mark profile exhausted after %v: %w", callErr, markErr)
|
||||
return llm.Message{}, "", total, fmt.Errorf("mark profile exhausted after %v: %w", callErr, markErr)
|
||||
}
|
||||
if transition.Advanced && !transition.Stale && hooks.transition != nil {
|
||||
hooks.transition(selection, transition, callErr)
|
||||
}
|
||||
if !transition.Stale && transition.NextProfileID == nil {
|
||||
return llm.Message{}, "", llm.Usage{}, &taskLLMError{taskID: taskID, chainExhausted: transition.ChainExhausted, cause: callErr}
|
||||
return llm.Message{}, "", total, &taskLLMError{taskID: taskID, chainExhausted: transition.ChainExhausted, cause: callErr}
|
||||
}
|
||||
continue // replay the same logical request on the next profile
|
||||
}
|
||||
return llm.Message{}, "", llm.Usage{}, callErr
|
||||
return llm.Message{}, "", total, callErr
|
||||
}
|
||||
// profileID=0: the explicit chain was cleared; non-quota error: pass through. Neither changes the task failover state.
|
||||
if selection.profileID == 0 || !isQuotaExhaustedError(callErr) {
|
||||
return llm.Message{}, "", llm.Usage{}, callErr
|
||||
return llm.Message{}, "", total, callErr
|
||||
}
|
||||
transition, markErr := hooks.exhaust(selection, callErr)
|
||||
if markErr != nil {
|
||||
return llm.Message{}, "", llm.Usage{}, fmt.Errorf("mark profile quota exhausted after %v: %w", callErr, markErr)
|
||||
return llm.Message{}, "", total, fmt.Errorf("mark profile quota exhausted after %v: %w", callErr, markErr)
|
||||
}
|
||||
if transition.Advanced && !transition.Stale && hooks.transition != nil {
|
||||
hooks.transition(selection, transition, callErr)
|
||||
}
|
||||
if !transition.Stale && transition.NextProfileID == nil {
|
||||
return llm.Message{}, "", llm.Usage{}, &taskLLMError{taskID: taskID, chainExhausted: transition.ChainExhausted, cause: callErr}
|
||||
return llm.Message{}, "", total, &taskLLMError{taskID: taskID, chainExhausted: transition.ChainExhausted, cause: callErr}
|
||||
}
|
||||
// No output has been delivered to the caller, so replaying the same logical request on the next profile is safe.
|
||||
}
|
||||
@@ -397,6 +402,35 @@ func streamTaskLLM(ctx context.Context, taskID string, req llm.CompletionRequest
|
||||
committed := false
|
||||
var pending []llm.StreamEvent
|
||||
var streamErr error
|
||||
// The caller's llm.Accumulator sums usage across the whole logical call
|
||||
// (every attempt and profile): message_start is added, message_delta is
|
||||
// folded with MergeEcho. Folding one attempt's echo against the sum of
|
||||
// earlier attempts would lose input/cache tokens, so each attempt's echoes
|
||||
// are normalized here against that attempt alone and only the increment
|
||||
// is passed on, as a message_start. The recorder below already billed the
|
||||
// physical attempt; this only shapes what the caller adds up.
|
||||
var attemptUsage llm.Usage
|
||||
forward := func(event llm.StreamEvent) bool {
|
||||
switch event.Type {
|
||||
case llm.SEMessageStart:
|
||||
attemptUsage.Add(event.Usage)
|
||||
case llm.SEMessageDelta:
|
||||
next := attemptUsage
|
||||
next.MergeEcho(event.Usage)
|
||||
inc := llm.Usage{
|
||||
InputTokens: next.InputTokens - attemptUsage.InputTokens,
|
||||
OutputTokens: next.OutputTokens - attemptUsage.OutputTokens,
|
||||
CacheReadTokens: next.CacheReadTokens - attemptUsage.CacheReadTokens,
|
||||
CacheWriteTokens: next.CacheWriteTokens - attemptUsage.CacheWriteTokens,
|
||||
}
|
||||
attemptUsage = next
|
||||
if inc != (llm.Usage{}) && !yield(llm.StreamEvent{Type: llm.SEMessageStart, Usage: inc}, nil) {
|
||||
return false
|
||||
}
|
||||
event.Usage = llm.Usage{}
|
||||
}
|
||||
return yield(event, nil)
|
||||
}
|
||||
retries, backoffOf := sameProviderRetryPolicy(selection.retry)
|
||||
// Same-provider safe-window retry: a transient failure before committed
|
||||
// (nothing delivered to the caller yet) can be replayed verbatim without
|
||||
@@ -405,6 +439,7 @@ func streamTaskLLM(ctx context.Context, taskID string, req llm.CompletionRequest
|
||||
committed = false
|
||||
pending = nil
|
||||
streamErr = nil
|
||||
attemptUsage = llm.Usage{}
|
||||
for event, err := range selection.provider.Stream(ctx, req) {
|
||||
if err != nil {
|
||||
streamErr = err
|
||||
@@ -416,24 +451,34 @@ func streamTaskLLM(ctx context.Context, taskID string, req llm.CompletionRequest
|
||||
}
|
||||
if !committed {
|
||||
for _, buffered := range pending {
|
||||
if !yield(buffered, nil) {
|
||||
if !forward(buffered) {
|
||||
return
|
||||
}
|
||||
}
|
||||
pending = nil
|
||||
committed = true
|
||||
}
|
||||
if !yield(event, nil) {
|
||||
if !forward(event) {
|
||||
return
|
||||
}
|
||||
}
|
||||
if streamErr != nil && !committed {
|
||||
// A failed uncommitted attempt delivers none of its events (they are
|
||||
// never replayed), but the usage it reported survives as one
|
||||
// message_start so the logical call still counts it.
|
||||
if u := pendingUsage(pending); u != (llm.Usage{}) && !yield(llm.StreamEvent{Type: llm.SEMessageStart, Usage: u}, nil) {
|
||||
return
|
||||
}
|
||||
pending = nil
|
||||
}
|
||||
if streamErr != nil && !committed && ctx.Err() == nil &&
|
||||
attempt < retries && isRetryableStreamError(streamErr) {
|
||||
backoff := backoffOf(attempt)
|
||||
log.Printf("[task-llm] task %s pre-commit stream failure, retrying on the same provider after %v (%d/%d): %v",
|
||||
taskID, backoff, attempt+1, retries, streamErr)
|
||||
if sleepCtx(ctx, backoff) {
|
||||
break // ctx cancelled during backoff → stop retrying
|
||||
yield(llm.StreamEvent{}, ctx.Err()) // cancelled during backoff → stop
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
@@ -442,7 +487,7 @@ func streamTaskLLM(ctx context.Context, taskID string, req llm.CompletionRequest
|
||||
if streamErr == nil {
|
||||
resetLLMFailStreak(taskID, selection.profileID)
|
||||
for _, buffered := range pending {
|
||||
if !yield(buffered, nil) {
|
||||
if !forward(buffered) {
|
||||
return
|
||||
}
|
||||
}
|
||||
@@ -474,11 +519,6 @@ func streamTaskLLM(ctx context.Context, taskID string, req llm.CompletionRequest
|
||||
// bundle was still in use. Agent/global fallback errors follow the legacy
|
||||
// behavior and never mutate task failover state.
|
||||
if selection.profileID == 0 || !isQuotaExhaustedError(streamErr) {
|
||||
for _, buffered := range pending {
|
||||
if !yield(buffered, nil) {
|
||||
return
|
||||
}
|
||||
}
|
||||
yield(llm.StreamEvent{}, streamErr)
|
||||
return
|
||||
}
|
||||
@@ -508,9 +548,25 @@ func streamTaskLLM(ctx context.Context, taskID string, req llm.CompletionRequest
|
||||
}
|
||||
}
|
||||
|
||||
// pendingUsage is the usage the buffered (uncommitted) events of one attempt
|
||||
// reported. Only message_start can be buffered: message_delta commits.
|
||||
func pendingUsage(pending []llm.StreamEvent) llm.Usage {
|
||||
var u llm.Usage
|
||||
for _, ev := range pending {
|
||||
if ev.Type == llm.SEMessageStart {
|
||||
u.Add(ev.Usage)
|
||||
}
|
||||
}
|
||||
return u
|
||||
}
|
||||
|
||||
// streamEventCommitsOutput reports whether delivering event makes the attempt
|
||||
// unreplayable. Signed and redacted thinking and tool-use signatures are opaque
|
||||
// output the model bound to this exact response, so they commit like text.
|
||||
func streamEventCommitsOutput(event llm.StreamEvent) bool {
|
||||
switch event.Type {
|
||||
case llm.SETextDelta, llm.SEThinkingDelta, llm.SEToolInputJSON:
|
||||
case llm.SETextDelta, llm.SEThinkingDelta, llm.SEToolInputJSON,
|
||||
llm.SEThinkingSignature, llm.SERedactedThinking, llm.SEToolUseSignature:
|
||||
return event.Text != ""
|
||||
case llm.SEToolUseStart, llm.SEMessageDelta, llm.SEMessageStop:
|
||||
return true
|
||||
|
||||
@@ -77,4 +77,3 @@ func TestEscalationAdvancesAfterConsecutiveFailures(t *testing.T) {
|
||||
t.Fatalf("exhaust calls=%d want 2", exhaustCalls)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,286 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"iter"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/llmrec"
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
)
|
||||
|
||||
// Usage reported by one failed attempt, and by a successful one. A stream
|
||||
// reports the input side in message_start and echoes it, with the output, in
|
||||
// message_delta; the atomic path returns the attempt's normalized usage.
|
||||
var (
|
||||
failedAttemptUsage = llm.Usage{InputTokens: 11, OutputTokens: 2, CacheReadTokens: 3, CacheWriteTokens: 2}
|
||||
okStartUsage = llm.Usage{InputTokens: 13, CacheReadTokens: 3, CacheWriteTokens: 2}
|
||||
okDeltaUsage = llm.Usage{InputTokens: 17, OutputTokens: 5, CacheReadTokens: 4, CacheWriteTokens: 2}
|
||||
okAttemptUsage = llm.Usage{InputTokens: 17, OutputTokens: 5, CacheReadTokens: 4, CacheWriteTokens: 2}
|
||||
)
|
||||
|
||||
// usageAttemptProvider fails its first fail attempts (always, when fail < 0)
|
||||
// after reporting failedAttemptUsage, then answers "done".
|
||||
type usageAttemptProvider struct {
|
||||
fail int
|
||||
failErr error
|
||||
onFail func()
|
||||
// extra events a failed stream attempt emits before its error.
|
||||
failEvents []llm.StreamEvent
|
||||
calls int
|
||||
}
|
||||
|
||||
func (p *usageAttemptProvider) failing() bool {
|
||||
p.calls++
|
||||
failed := p.fail < 0 || p.calls <= p.fail
|
||||
if failed && p.onFail != nil {
|
||||
p.onFail()
|
||||
}
|
||||
return failed
|
||||
}
|
||||
|
||||
func (p *usageAttemptProvider) Complete(context.Context, llm.CompletionRequest) (llm.Message, string, llm.Usage, error) {
|
||||
if p.failing() {
|
||||
return llm.Message{}, "", failedAttemptUsage, p.failErr
|
||||
}
|
||||
return llm.Message{Role: llm.RoleAssistant, Content: []llm.ContentBlock{llm.TextBlock("done")}}, "end_turn", okAttemptUsage, nil
|
||||
}
|
||||
|
||||
func (p *usageAttemptProvider) Stream(context.Context, llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
|
||||
return func(yield func(llm.StreamEvent, error) bool) {
|
||||
if p.failing() {
|
||||
events := append([]llm.StreamEvent{{Type: llm.SEMessageStart, Usage: failedAttemptUsage}}, p.failEvents...)
|
||||
for _, ev := range events {
|
||||
if !yield(ev, nil) {
|
||||
return
|
||||
}
|
||||
}
|
||||
yield(llm.StreamEvent{}, p.failErr)
|
||||
return
|
||||
}
|
||||
for _, ev := range []llm.StreamEvent{
|
||||
{Type: llm.SEMessageStart, Usage: okStartUsage},
|
||||
{Type: llm.SETextDelta, Text: "done"},
|
||||
{Type: llm.SEMessageDelta, StopReason: "end_turn", Usage: okDeltaUsage},
|
||||
{Type: llm.SEMessageStop},
|
||||
} {
|
||||
if !yield(ev, nil) {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// logicalResult is what the SDK sees for one logical call.
|
||||
type logicalResult struct {
|
||||
usage llm.Usage
|
||||
text string
|
||||
events []llm.StreamEvent
|
||||
err error
|
||||
}
|
||||
|
||||
func runLogicalCall(ctx context.Context, stream bool, taskID string, hooks taskLLMStreamHooks) logicalResult {
|
||||
if !stream {
|
||||
msg, _, usage, err := completeTaskLLM(ctx, taskID, llm.CompletionRequest{}, hooks)
|
||||
return logicalResult{usage: usage, text: msg.Text(), err: err}
|
||||
}
|
||||
var r logicalResult
|
||||
acc := llm.NewAccumulator()
|
||||
for ev, err := range streamTaskLLM(ctx, taskID, llm.CompletionRequest{}, hooks) {
|
||||
if err != nil {
|
||||
r.err = err
|
||||
break
|
||||
}
|
||||
r.events = append(r.events, ev)
|
||||
acc.Add(ev)
|
||||
}
|
||||
r.usage, r.text = acc.Usage, acc.Message().Text()
|
||||
return r
|
||||
}
|
||||
|
||||
func singleProfileHooks(p llm.Provider, retry agent.RetryConfig) taskLLMStreamHooks {
|
||||
return taskLLMStreamHooks{
|
||||
current: func(context.Context) (taskLLMSelection, error) {
|
||||
return taskLLMSelection{provider: p, retry: retry}, nil
|
||||
},
|
||||
exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) {
|
||||
return db.TaskLLMTransition{}, nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func sumUsage(us ...llm.Usage) llm.Usage {
|
||||
var t llm.Usage
|
||||
for _, u := range us {
|
||||
t.Add(u)
|
||||
}
|
||||
return t
|
||||
}
|
||||
|
||||
var usageModes = []struct {
|
||||
name string
|
||||
stream bool
|
||||
}{{"atomic", false}, {"stream", true}}
|
||||
|
||||
func TestTaskLLMRetrySumsAttemptUsage(t *testing.T) {
|
||||
for _, m := range usageModes {
|
||||
t.Run(m.name, func(t *testing.T) {
|
||||
p := &usageAttemptProvider{fail: 1, failErr: errors.New("connection reset by peer")}
|
||||
r := runLogicalCall(t.Context(), m.stream, "p12-retry", singleProfileHooks(p, agent.RetryConfig{StreamAttempts: 1, StreamInterval: time.Millisecond}))
|
||||
if r.err != nil {
|
||||
t.Fatal(r.err)
|
||||
}
|
||||
want := sumUsage(failedAttemptUsage, okAttemptUsage) // {28 7 7 4}
|
||||
if r.usage != want || p.calls != 2 || r.text != "done" {
|
||||
t.Fatalf("usage=%+v want %+v, calls=%d, text=%q", r.usage, want, p.calls, r.text)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskLLMExhaustedRetriesReturnSummedUsage(t *testing.T) {
|
||||
for _, m := range usageModes {
|
||||
t.Run(m.name, func(t *testing.T) {
|
||||
p := &usageAttemptProvider{fail: -1, failErr: errors.New("connection reset by peer")}
|
||||
r := runLogicalCall(t.Context(), m.stream, "p12-exhausted", singleProfileHooks(p, agent.RetryConfig{StreamAttempts: 1, StreamInterval: time.Millisecond}))
|
||||
if r.err == nil || r.err.Error() != "connection reset by peer" {
|
||||
t.Fatalf("err=%v", r.err)
|
||||
}
|
||||
want := sumUsage(failedAttemptUsage, failedAttemptUsage)
|
||||
if r.usage != want || p.calls != 2 {
|
||||
t.Fatalf("usage=%+v want %+v, calls=%d", r.usage, want, p.calls)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskLLMProfileFailoverSumsUsage(t *testing.T) {
|
||||
for _, m := range usageModes {
|
||||
t.Run(m.name, func(t *testing.T) {
|
||||
first := &usageAttemptProvider{fail: -1, failErr: errors.New("openai: status 402: insufficient_quota")}
|
||||
second := &usageAttemptProvider{}
|
||||
active := int64(11)
|
||||
hooks := taskLLMStreamHooks{
|
||||
current: func(context.Context) (taskLLMSelection, error) {
|
||||
if active == 11 {
|
||||
return taskLLMSelection{profileID: 11, provider: first}, nil
|
||||
}
|
||||
return taskLLMSelection{profileID: 22, provider: second}, nil
|
||||
},
|
||||
exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) {
|
||||
active = 22
|
||||
next := active
|
||||
return db.TaskLLMTransition{PreviousProfileID: 11, NextProfileID: &next, Advanced: true}, nil
|
||||
},
|
||||
}
|
||||
r := runLogicalCall(t.Context(), m.stream, "p12-failover-"+m.name, hooks)
|
||||
if r.err != nil {
|
||||
t.Fatal(r.err)
|
||||
}
|
||||
want := sumUsage(failedAttemptUsage, okAttemptUsage)
|
||||
if r.usage != want || first.calls != 1 || second.calls != 1 || r.text != "done" {
|
||||
t.Fatalf("usage=%+v want %+v, calls first=%d second=%d, text=%q", r.usage, want, first.calls, second.calls, r.text)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskLLMCancelDuringBackoffStops(t *testing.T) {
|
||||
for _, m := range usageModes {
|
||||
t.Run(m.name, func(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(t.Context())
|
||||
defer cancel()
|
||||
p := &usageAttemptProvider{fail: -1, failErr: errors.New("connection reset by peer")}
|
||||
// Cancel once the retry is waiting out its (long) backoff.
|
||||
p.onFail = func() { time.AfterFunc(20*time.Millisecond, cancel) }
|
||||
r := runLogicalCall(ctx, m.stream, "p12-cancel", singleProfileHooks(p, agent.RetryConfig{StreamAttempts: 2, StreamInterval: time.Minute}))
|
||||
if !errors.Is(r.err, context.Canceled) {
|
||||
t.Fatalf("err=%v, want context.Canceled", r.err)
|
||||
}
|
||||
if p.calls != 1 || r.usage != failedAttemptUsage {
|
||||
t.Fatalf("calls=%d usage=%+v", p.calls, r.usage)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskLLMUncommittedStreamAttemptKeepsUsageWithoutReplay(t *testing.T) {
|
||||
p := &usageAttemptProvider{
|
||||
fail: 1,
|
||||
failErr: errors.New("anthropic: overloaded_error"),
|
||||
// Buffered, non-committing events of the failed attempt.
|
||||
failEvents: []llm.StreamEvent{{Type: llm.SEThinkingDelta}, {Type: llm.SEToolInputJSON}},
|
||||
}
|
||||
r := runLogicalCall(t.Context(), true, "p12-uncommitted", singleProfileHooks(p, agent.RetryConfig{StreamAttempts: 1, StreamInterval: time.Millisecond}))
|
||||
if r.err != nil {
|
||||
t.Fatal(r.err)
|
||||
}
|
||||
if want := sumUsage(failedAttemptUsage, okAttemptUsage); r.usage != want {
|
||||
t.Fatalf("usage=%+v want %+v", r.usage, want)
|
||||
}
|
||||
if r.text != "done" {
|
||||
t.Fatalf("text=%q, want the second attempt's only", r.text)
|
||||
}
|
||||
// The failed attempt contributes one usage-only message_start, nothing else.
|
||||
if len(r.events) == 0 || r.events[0].Type != llm.SEMessageStart || r.events[0].Usage != failedAttemptUsage {
|
||||
t.Fatalf("first event %+v", r.events)
|
||||
}
|
||||
for _, ev := range r.events {
|
||||
if ev.Type == llm.SEThinkingDelta || ev.Type == llm.SEToolInputJSON {
|
||||
t.Fatalf("buffered event of the failed attempt was delivered: %+v", r.events)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskLLMOpaqueThinkingCommitsStream(t *testing.T) {
|
||||
defer withZeroRetryBackoff()()
|
||||
for _, kind := range []llm.StreamEventType{llm.SEThinkingSignature, llm.SERedactedThinking, llm.SEToolUseSignature} {
|
||||
t.Run(string(kind), func(t *testing.T) {
|
||||
if !streamEventCommitsOutput(llm.StreamEvent{Type: kind, Text: "opaque"}) {
|
||||
t.Fatalf("%s does not commit", kind)
|
||||
}
|
||||
// A retryable failure after the opaque block is surfaced, not replayed.
|
||||
p := &scriptedLLMProvider{events: []llm.StreamEvent{{Type: kind, Text: "opaque"}}, err: errors.New("connection reset by peer")}
|
||||
events, err := collectTaskLLMStream(streamTaskLLM(t.Context(), "p12-opaque", llm.CompletionRequest{}, singleProfileHooks(p, agent.RetryConfig{})))
|
||||
if err == nil || p.calls != 1 || len(events) != 1 || events[0].Type != kind {
|
||||
t.Fatalf("calls=%d events=%+v err=%v", p.calls, events, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// The recorder below the task router writes one llm_usage row per physical
|
||||
// attempt; summing the attempts for the logical call adds no row.
|
||||
func TestTaskLLMRetryKeepsOneLedgerRowPerPhysicalCall(t *testing.T) {
|
||||
m, err := NewManager(t.TempDir(), "")
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
defer m.Close()
|
||||
pg := m.PG()
|
||||
for _, mode := range usageModes {
|
||||
t.Run(mode.name, func(t *testing.T) {
|
||||
taskID := fmt.Sprintf("p12-ledger-%s-%d", mode.name, time.Now().UnixNano())
|
||||
t.Cleanup(func() { _, _ = pg.Exec(`DELETE FROM llm_usage WHERE task_id=$1`, taskID) })
|
||||
p := &usageAttemptProvider{fail: 1, failErr: errors.New("connection reset by peer")}
|
||||
rec := llmrec.Wrap(p, pg, "test-model", "test-profile", "", "", func() bool { return false })
|
||||
ctx := llmrec.WithTaskID(t.Context(), taskID)
|
||||
r := runLogicalCall(ctx, mode.stream, taskID, singleProfileHooks(rec, agent.RetryConfig{StreamAttempts: 1, StreamInterval: time.Millisecond}))
|
||||
if r.err != nil {
|
||||
t.Fatal(r.err)
|
||||
}
|
||||
var rows, failed int
|
||||
if err := pg.QueryRow(`SELECT count(*), count(*) FILTER (WHERE status='error') FROM llm_usage WHERE task_id=$1`, taskID).Scan(&rows, &failed); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if rows != 2 || failed != 1 || p.calls != 2 {
|
||||
t.Fatalf("llm_usage rows=%d (error %d), physical calls=%d; want 2 rows, 1 error", rows, failed, p.calls)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -5,9 +5,9 @@ import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
)
|
||||
|
||||
func names(tools []actool.CoreTool) map[string]actool.CoreTool {
|
||||
|
||||
@@ -32,8 +32,11 @@ func (s SideQuestionService) Answer(ctx context.Context, req llm.CompletionReque
|
||||
out.Text += ev.Text
|
||||
case llm.SEToolUseStart:
|
||||
out.ToolUse = true
|
||||
case llm.SEMessageStart, llm.SEMessageDelta:
|
||||
case llm.SEMessageStart:
|
||||
out.Usage.Add(ev.Usage)
|
||||
case llm.SEMessageDelta:
|
||||
// message_delta may echo message_start's input side: MergeEcho, not Add.
|
||||
out.Usage.MergeEcho(ev.Usage)
|
||||
case llm.SEMessageStop:
|
||||
complete = true
|
||||
}
|
||||
|
||||
@@ -0,0 +1,52 @@
|
||||
package sidequestion
|
||||
|
||||
import (
|
||||
"context"
|
||||
"iter"
|
||||
"testing"
|
||||
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
)
|
||||
|
||||
// A streamed side answer counts message_delta's echo of the input side once.
|
||||
func TestAnswerStreamUsageEchoCountedOnce(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
start, delta llm.Usage
|
||||
want llm.Usage
|
||||
}{
|
||||
{"full_echo",
|
||||
llm.Usage{InputTokens: 23, CacheReadTokens: 5, CacheWriteTokens: 3},
|
||||
llm.Usage{InputTokens: 23, CacheReadTokens: 5, CacheWriteTokens: 3, OutputTokens: 5},
|
||||
llm.Usage{InputTokens: 23, OutputTokens: 5, CacheReadTokens: 5, CacheWriteTokens: 3}},
|
||||
{"partial_echo",
|
||||
llm.Usage{InputTokens: 740, CacheReadTokens: 200, CacheWriteTokens: 40},
|
||||
llm.Usage{InputTokens: 500, OutputTokens: 77},
|
||||
llm.Usage{InputTokens: 740, OutputTokens: 77, CacheReadTokens: 200, CacheWriteTokens: 40}},
|
||||
{"no_echo",
|
||||
llm.Usage{InputTokens: 23, CacheReadTokens: 5, CacheWriteTokens: 3},
|
||||
llm.Usage{OutputTokens: 5},
|
||||
llm.Usage{InputTokens: 23, OutputTokens: 5, CacheReadTokens: 5, CacheWriteTokens: 3}},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
p := fakeProvider{stream: func(context.Context, llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
|
||||
return func(yield func(llm.StreamEvent, error) bool) {
|
||||
for _, ev := range []llm.StreamEvent{
|
||||
{Type: llm.SEMessageStart, Usage: tc.start},
|
||||
{Type: llm.SETextDelta, Text: "answer"},
|
||||
{Type: llm.SEMessageDelta, StopReason: "end_turn", Usage: tc.delta},
|
||||
{Type: llm.SEMessageStop},
|
||||
} {
|
||||
if !yield(ev, nil) {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}}
|
||||
out, err := (SideQuestionService{Provider: p}).Answer(t.Context(), llm.CompletionRequest{}, true, nil)
|
||||
if err != nil || out.Text != "answer" || out.Usage != tc.want {
|
||||
t.Fatalf("answer=%+v err=%v, want usage %+v", out, err, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package traffic
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/egress"
|
||||
mproxy "github.com/lqqyt2423/go-mitmproxy/proxy"
|
||||
)
|
||||
|
||||
// paceWaitMax bounds the pacing wait of one request. Past it the request is
|
||||
// sent unpaced: pacing never blocks a command indefinitely.
|
||||
const paceWaitMax = 30 * time.Second
|
||||
|
||||
// Requestheaders runs for every request the proxy sees: the CONNECT that opens
|
||||
// a tunnel, each request inside an intercepted tunnel, and each plain-HTTP
|
||||
// request. The recording proxy takes no credentials, so a Proxy-Authorization
|
||||
// header is only ever meant for it (clients send the proxy URL's userinfo
|
||||
// there): it is dropped here, before anything is forwarded, and never reaches
|
||||
// the target or the recording. When it carries the pace marker
|
||||
// (egress.PaceUser), the client connection is marked; every request on it,
|
||||
// including those inside its tunnel, is then paced by the current noise
|
||||
// profile before it goes upstream. Unmarked connections are not paced here.
|
||||
func (s *sink) Requestheaders(f *mproxy.Flow) {
|
||||
if f.Request == nil || f.ConnContext == nil || f.ConnContext.ClientConn == nil {
|
||||
return
|
||||
}
|
||||
id := f.ConnContext.ClientConn.Id
|
||||
if egress.IsPaceMarker(f.Request.Header.Get("Proxy-Authorization")) {
|
||||
s.paced.Store(id, struct{}{})
|
||||
}
|
||||
f.Request.Header.Del("Proxy-Authorization")
|
||||
if f.Request.Method == http.MethodConnect || f.Request.URL == nil {
|
||||
return
|
||||
}
|
||||
if _, marked := s.paced.Load(id); marked {
|
||||
pace(f)
|
||||
}
|
||||
}
|
||||
|
||||
// ClientDisconnected forgets a closed connection's mark.
|
||||
func (s *sink) ClientDisconnected(c *mproxy.ClientConn) {
|
||||
if c != nil {
|
||||
s.paced.Delete(c.Id)
|
||||
}
|
||||
}
|
||||
|
||||
// pace waits for the request's host under the current noise profile, bounded
|
||||
// by paceWaitMax and by the request's own context. Off returns at once.
|
||||
func pace(f *mproxy.Flow) {
|
||||
ctx := context.Background()
|
||||
if raw := f.Request.Raw(); raw != nil {
|
||||
ctx = raw.Context()
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(ctx, paceWaitMax)
|
||||
defer cancel()
|
||||
host := egress.HostKey(f.Request.URL)
|
||||
if err := egress.Wait(ctx, host); err != nil {
|
||||
log.Printf("[traffic] pacing a Bash request to %s stopped waiting (%v); sending it unpaced", host, err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,262 @@
|
||||
package traffic
|
||||
|
||||
import (
|
||||
"crypto/tls"
|
||||
"crypto/x509"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"os"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/egress"
|
||||
mproxy "github.com/lqqyt2423/go-mitmproxy/proxy"
|
||||
)
|
||||
|
||||
// startProxy runs a real recording proxy on a free loopback port and returns it
|
||||
// with its sink and URL. extra addons run after the recording sink.
|
||||
func startProxy(t *testing.T, extra ...mproxy.Addon) (*Traffic, *sink, string) {
|
||||
t.Helper()
|
||||
ln, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
addr := ln.Addr().String()
|
||||
ln.Close()
|
||||
tr, err := Open(t.TempDir(), addr)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s := tr.proxy.Addons[0].(*sink)
|
||||
for _, a := range extra {
|
||||
tr.proxy.AddAddon(a)
|
||||
}
|
||||
go func() { _ = tr.Start() }()
|
||||
t.Cleanup(func() {
|
||||
_ = tr.proxy.Close()
|
||||
_ = tr.Close()
|
||||
})
|
||||
deadline := time.Now().Add(5 * time.Second)
|
||||
for {
|
||||
c, err := net.Dial("tcp", addr)
|
||||
if err == nil {
|
||||
c.Close()
|
||||
break
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatalf("proxy did not listen on %s: %v", addr, err)
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
return tr, s, "http://" + addr
|
||||
}
|
||||
|
||||
// proxyClient is a curl-like client: it sends the proxy URL's userinfo as
|
||||
// Proxy-Authorization on CONNECT and on plain-HTTP requests, and trusts the CA.
|
||||
func proxyClient(t *testing.T, tr *Traffic, proxyURL string) *http.Client {
|
||||
t.Helper()
|
||||
pem, err := os.ReadFile(tr.CACertPath())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
pool := x509.NewCertPool()
|
||||
if !pool.AppendCertsFromPEM(pem) {
|
||||
t.Fatal("no CA in", tr.CACertPath())
|
||||
}
|
||||
u, _ := url.Parse(proxyURL)
|
||||
return &http.Client{Transport: &http.Transport{Proxy: http.ProxyURL(u), TLSClientConfig: &tls.Config{RootCAs: pool}}}
|
||||
}
|
||||
|
||||
// probe records the ClientConn of each request the proxy sees, after the sink.
|
||||
type probe struct {
|
||||
mproxy.BaseAddon
|
||||
mu sync.Mutex
|
||||
conn map[string]string // "METHOD url" -> ClientConn.Id
|
||||
}
|
||||
|
||||
func (p *probe) Requestheaders(f *mproxy.Flow) {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
if p.conn == nil {
|
||||
p.conn = map[string]string{}
|
||||
}
|
||||
p.conn[f.Request.Method+" "+f.Request.URL.String()] = f.ConnContext.ClientConn.Id.String()
|
||||
}
|
||||
|
||||
func get(t *testing.T, c *http.Client, u string) {
|
||||
t.Helper()
|
||||
resp, err := c.Get(u)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, _ = io.Copy(io.Discard, resp.Body)
|
||||
resp.Body.Close()
|
||||
if resp.StatusCode != 200 {
|
||||
t.Fatalf("GET %s: %d", u, resp.StatusCode)
|
||||
}
|
||||
}
|
||||
|
||||
func markedCount(s *sink) int {
|
||||
n := 0
|
||||
s.paced.Range(func(any, any) bool { n++; return true })
|
||||
return n
|
||||
}
|
||||
|
||||
func hostRPM(host string) int {
|
||||
for _, h := range egress.Current().Snapshot() {
|
||||
if h.Host == host {
|
||||
return h.RPM
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// A marked client (the Bash proxy URL) is paced and counted, through a CONNECT
|
||||
// tunnel and on plain HTTP; an unmarked one (WebFetch, custom tools, enrich)
|
||||
// is not paced here. No Proxy-Authorization reaches a target or the recording,
|
||||
// every exchange is recorded, and a closed connection loses its mark.
|
||||
func TestBashPacingThroughProxy(t *testing.T) {
|
||||
prev := egress.CurrentProfile()
|
||||
t.Cleanup(func() { egress.SetProfile(prev) })
|
||||
// 2 req/s, burst 1, no jitter: from the second request on, each waits 500 ms.
|
||||
egress.SetProfile(egress.Profile{Name: "test", RPS: 2, Burst: 1, MaxPerHost: 1, MaxGlobal: 8})
|
||||
|
||||
var mu sync.Mutex
|
||||
leaked := map[string]string{}
|
||||
handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if v := r.Header.Get("Proxy-Authorization"); v != "" {
|
||||
mu.Lock()
|
||||
leaked[r.URL.Path] = v
|
||||
mu.Unlock()
|
||||
}
|
||||
_, _ = io.WriteString(w, "ok")
|
||||
})
|
||||
marked := httptest.NewServer(handler)
|
||||
defer marked.Close()
|
||||
unmarked := httptest.NewServer(handler)
|
||||
defer unmarked.Close()
|
||||
tlsSrv := httptest.NewTLSServer(handler)
|
||||
defer tlsSrv.Close()
|
||||
// The MITM leaf for a bare IP has no IP SAN: address the TLS target by name.
|
||||
tlsURL := strings.Replace(tlsSrv.URL, "127.0.0.1", "localhost", 1)
|
||||
|
||||
p := &probe{}
|
||||
tr, s, proxyURL := startProxy(t, p)
|
||||
bash := proxyClient(t, tr, egress.MarkProxyURL(proxyURL))
|
||||
plain := proxyClient(t, tr, proxyURL)
|
||||
|
||||
start := time.Now()
|
||||
for range 3 {
|
||||
get(t, bash, marked.URL+"/bash")
|
||||
}
|
||||
if d := time.Since(start); d < 900*time.Millisecond {
|
||||
t.Fatalf("3 marked requests to one host took %v, want >= 1s of pacing", d)
|
||||
}
|
||||
if n := hostRPM(strings.TrimPrefix(marked.URL, "http://")); n != 3 {
|
||||
t.Fatalf("egress counter for the marked host = %d, want 3", n)
|
||||
}
|
||||
for range 3 {
|
||||
get(t, plain, unmarked.URL+"/plain")
|
||||
}
|
||||
if n := hostRPM(strings.TrimPrefix(unmarked.URL, "http://")); n != 0 {
|
||||
t.Fatalf("unmarked requests were paced by the proxy: counter %d, want 0", n)
|
||||
}
|
||||
get(t, bash, tlsURL+"/tls1")
|
||||
get(t, bash, tlsURL+"/tls2")
|
||||
tlsHost := strings.TrimPrefix(tlsURL, "https://")
|
||||
if n := hostRPM(tlsHost); n != 2 {
|
||||
t.Fatalf("egress counter for the tunnelled host = %d, want 2", n)
|
||||
}
|
||||
|
||||
p.mu.Lock()
|
||||
connectID := p.conn["CONNECT //"+tlsHost]
|
||||
inner1, inner2 := p.conn["GET "+tlsURL+"/tls1"], p.conn["GET "+tlsURL+"/tls2"]
|
||||
p.mu.Unlock()
|
||||
if connectID == "" || inner1 != connectID || inner2 != connectID {
|
||||
t.Fatalf("tunnel flows do not share the CONNECT's connection: connect=%q inner=%q,%q", connectID, inner1, inner2)
|
||||
}
|
||||
mu.Lock()
|
||||
if len(leaked) != 0 {
|
||||
t.Fatalf("Proxy-Authorization reached the target: %v", leaked)
|
||||
}
|
||||
mu.Unlock()
|
||||
|
||||
rows, err := tr.Search("", 0, 100)
|
||||
if err != nil || len(rows) != 8 {
|
||||
t.Fatalf("recorded %d exchanges (err %v), want 8", len(rows), err)
|
||||
}
|
||||
for _, r := range rows {
|
||||
req, _, err := tr.Get(r.ID)
|
||||
if err != nil || strings.Contains(strings.ToLower(req), "proxy-authorization") {
|
||||
t.Fatalf("recorded request %s keeps Proxy-Authorization (err %v):\n%s", r.ID, err, req)
|
||||
}
|
||||
}
|
||||
|
||||
if markedCount(s) == 0 {
|
||||
t.Fatal("no connection marked")
|
||||
}
|
||||
bash.CloseIdleConnections()
|
||||
plain.CloseIdleConnections()
|
||||
deadline := time.Now().Add(5 * time.Second)
|
||||
for markedCount(s) != 0 {
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatalf("%d marks left after the clients closed", markedCount(s))
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
}
|
||||
|
||||
// Off: a marked request is neither paced nor counted.
|
||||
func TestBashPacingOff(t *testing.T) {
|
||||
prev := egress.CurrentProfile()
|
||||
t.Cleanup(func() { egress.SetProfile(prev) })
|
||||
egress.SetProfile(egress.Profile{Name: "off"})
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
|
||||
defer srv.Close()
|
||||
tr, _, proxyURL := startProxy(t)
|
||||
bash := proxyClient(t, tr, egress.MarkProxyURL(proxyURL))
|
||||
for range 3 {
|
||||
get(t, bash, srv.URL+"/")
|
||||
}
|
||||
if snap := egress.Current().Snapshot(); len(snap) != 0 {
|
||||
t.Fatalf("off counted marked requests: %v", snap)
|
||||
}
|
||||
}
|
||||
|
||||
// The sink drops Proxy-Authorization from every request and marks only on the
|
||||
// pace marker; a disconnect clears the mark.
|
||||
func TestSinkStripsAndMarks(t *testing.T) {
|
||||
prev := egress.CurrentProfile()
|
||||
t.Cleanup(func() { egress.SetProfile(prev) })
|
||||
egress.SetProfile(egress.Profile{Name: "off"})
|
||||
s := &sink{}
|
||||
flow := func(method, auth string, conn *mproxy.ClientConn) *mproxy.Flow {
|
||||
u, _ := url.Parse("http://a.example/")
|
||||
h := http.Header{}
|
||||
if auth != "" {
|
||||
h.Set("Proxy-Authorization", auth)
|
||||
}
|
||||
return &mproxy.Flow{Request: &mproxy.Request{Method: method, URL: u, Header: h}, ConnContext: &mproxy.ConnContext{ClientConn: conn}}
|
||||
}
|
||||
f := flow(http.MethodGet, "Basic dXNlcjpwYXNz", &mproxy.ClientConn{}) // user:pass, not the marker
|
||||
s.Requestheaders(f)
|
||||
if f.Request.Header.Get("Proxy-Authorization") != "" || markedCount(s) != 0 {
|
||||
t.Fatalf("foreign credentials: header %q, marks %d", f.Request.Header.Get("Proxy-Authorization"), markedCount(s))
|
||||
}
|
||||
conn := &mproxy.ClientConn{}
|
||||
conn.Id[0] = 1
|
||||
f = flow(http.MethodConnect, "Basic YXJ0ZXgtcGFjZTphcnRleC1wYWNl", conn) // artex-pace:artex-pace
|
||||
s.Requestheaders(f)
|
||||
if f.Request.Header.Get("Proxy-Authorization") != "" || markedCount(s) != 1 {
|
||||
t.Fatalf("marker on CONNECT: header %q, marks %d", f.Request.Header.Get("Proxy-Authorization"), markedCount(s))
|
||||
}
|
||||
s.ClientDisconnected(conn)
|
||||
if markedCount(s) != 0 {
|
||||
t.Fatal("mark kept after disconnect")
|
||||
}
|
||||
}
|
||||
+4
-2
@@ -383,10 +383,12 @@ func (t *Traffic) stopping() bool {
|
||||
}
|
||||
func (t *Traffic) DB() *sql.DB { return t.db }
|
||||
|
||||
// sink is the go-mitmproxy addon that records completed exchanges.
|
||||
// sink is the go-mitmproxy addon that records completed exchanges, and paces
|
||||
// the requests of client connections marked for pacing (see pace.go).
|
||||
type sink struct {
|
||||
mproxy.BaseAddon
|
||||
t *Traffic
|
||||
t *Traffic
|
||||
paced sync.Map // ClientConn.Id of a marked connection -> struct{}
|
||||
}
|
||||
|
||||
func (s *sink) Response(f *mproxy.Flow) {
|
||||
|
||||
Generated
+588
-234
File diff suppressed because it is too large
Load Diff
+2
-3
@@ -21,7 +21,6 @@
|
||||
},
|
||||
"dependencies": {
|
||||
"@antv/g6": "^5.1.1",
|
||||
"@antv/graphin": "^3.0.5",
|
||||
"@base-ui/react": "^1.6.0",
|
||||
"@dnd-kit/core": "^6.3.1",
|
||||
"@dnd-kit/modifiers": "^9.0.0",
|
||||
@@ -40,7 +39,7 @@
|
||||
"geist": "^1.7.2",
|
||||
"input-otp": "^1.4.2",
|
||||
"lucide-react": "^1.22.0",
|
||||
"next": "^16.2.9",
|
||||
"next": "^16.4.6",
|
||||
"next-themes": "^0.4.6",
|
||||
"radix-ui": "^1.6.0",
|
||||
"react": "^19.2.7",
|
||||
@@ -52,7 +51,6 @@
|
||||
"react-resizable-panels": "^4.12.0",
|
||||
"recharts": "^3.8.0",
|
||||
"remark-gfm": "^4.0.1",
|
||||
"shadcn": "^4.12.0",
|
||||
"simple-icons": "^16.24.1",
|
||||
"sonner": "^2.0.7",
|
||||
"tailwind-merge": "^3.6.0",
|
||||
@@ -74,6 +72,7 @@
|
||||
"husky": "^9.1.7",
|
||||
"lint-staged": "^16.4.0",
|
||||
"postcss": "^8.5.16",
|
||||
"shadcn": "^4.12.0",
|
||||
"tailwindcss": "^4.1.5",
|
||||
"ts-node": "^10.9.2",
|
||||
"tw-animate-css": "^1.4.0",
|
||||
|
||||
@@ -0,0 +1,153 @@
|
||||
"use client";
|
||||
|
||||
import * as React from "react";
|
||||
|
||||
import { GaugeIcon, RefreshCwIcon } from "lucide-react";
|
||||
import { toast } from "sonner";
|
||||
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { Card, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card";
|
||||
import { Label } from "@/components/ui/label";
|
||||
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
|
||||
import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table";
|
||||
import { api } from "@/lib/api";
|
||||
import type { NoiseProfileName, NoiseStatus } from "@/lib/types";
|
||||
|
||||
// Mirrors egress.Profiles in egress/egress.go; the jitter of one send is
|
||||
// min(Jitter/RPS, 250 ms).
|
||||
const profiles: { value: NoiseProfileName; label: string; limits: string }[] = [
|
||||
{ value: "off", label: "Off", limits: "Off: requests and their timing are unchanged, and nothing is counted." },
|
||||
{
|
||||
value: "quiet",
|
||||
label: "Quiet",
|
||||
limits:
|
||||
"2 req/s per host (burst 4), 1 in flight per host, 8 overall, jitter up to 250 ms, identical GETs of a task cached 60 s, up to 3 retries on 429/503.",
|
||||
},
|
||||
{
|
||||
value: "balanced",
|
||||
label: "Balanced",
|
||||
limits:
|
||||
"5 req/s per host (burst 10), 2 in flight per host, 16 overall, jitter up to 40 ms, identical GETs of a task cached 30 s, up to 3 retries on 429/503.",
|
||||
},
|
||||
{
|
||||
value: "aggressive",
|
||||
label: "Aggressive",
|
||||
limits:
|
||||
"20 req/s per host (burst 40), 4 in flight per host, 32 overall, jitter up to 5 ms, identical GETs of a task cached 10 s, up to 2 retries on 429/503.",
|
||||
},
|
||||
];
|
||||
|
||||
// NoiseProfileCard sets the target noise profile and shows the per-host
|
||||
// counters of the current minute.
|
||||
export function NoiseProfileCard() {
|
||||
const [status, setStatus] = React.useState<NoiseStatus | null>(null);
|
||||
const [loading, setLoading] = React.useState(false);
|
||||
const [saving, setSaving] = React.useState(false);
|
||||
|
||||
const refresh = React.useCallback(() => {
|
||||
setLoading(true);
|
||||
api
|
||||
.noiseStatus()
|
||||
.then(setStatus)
|
||||
.catch((e) => toast.error(`Failed to load the noise counters: ${(e as Error).message}`))
|
||||
.finally(() => setLoading(false));
|
||||
}, []);
|
||||
React.useEffect(() => {
|
||||
refresh();
|
||||
}, [refresh]);
|
||||
|
||||
const save = (value: NoiseProfileName) => {
|
||||
setSaving(true);
|
||||
api
|
||||
.setSettings({ noise_profile: value })
|
||||
.then(() => {
|
||||
toast.success(`Noise profile set to ${value}`);
|
||||
refresh();
|
||||
})
|
||||
.catch((e) => toast.error(`Save failed: ${(e as Error).message}`))
|
||||
.finally(() => setSaving(false));
|
||||
};
|
||||
|
||||
const current = profiles.find((p) => p.value === status?.profile) ?? profiles[0];
|
||||
return (
|
||||
<Card className="mb-4 break-inside-avoid md:mb-6">
|
||||
<CardHeader>
|
||||
<CardTitle className="flex items-center gap-2 text-base">
|
||||
<GaugeIcon className="size-4" />
|
||||
Target noise
|
||||
</CardTitle>
|
||||
<CardDescription>
|
||||
Paces the requests ARTEX itself sends to targets (WebFetch, http custom tools, the enrich probe) and sets the
|
||||
low-and-slow guidance in agent prompts. With traffic capture on, Bash commands that use the proxy (curl,
|
||||
wget, python-requests…) are paced too, by rate and jitter only: the in-flight, cache and retry limits apply
|
||||
to ARTEX's own requests. Raw scanners such as nmap, and commands that bypass the proxy, are not paced; this
|
||||
is best effort, not a security boundary. This setting takes precedence over ARTEX_NOISE_PROFILE and applies
|
||||
to the next request.
|
||||
</CardDescription>
|
||||
</CardHeader>
|
||||
<CardContent className="flex flex-col gap-4">
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<Label htmlFor="noise-profile" className="text-sm font-normal text-muted-foreground">
|
||||
Profile
|
||||
</Label>
|
||||
<Select
|
||||
value={status?.profile ?? ""}
|
||||
disabled={!status || saving}
|
||||
onValueChange={(v) => save(v as NoiseProfileName)}
|
||||
>
|
||||
<SelectTrigger id="noise-profile" className="w-48 shrink-0">
|
||||
<SelectValue placeholder="Loading…" />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
{profiles.map((p) => (
|
||||
<SelectItem key={p.value} value={p.value}>
|
||||
{p.label}
|
||||
</SelectItem>
|
||||
))}
|
||||
</SelectContent>
|
||||
</Select>
|
||||
</div>
|
||||
<p className="text-muted-foreground text-xs">{current.limits}</p>
|
||||
<div className="flex items-center justify-between gap-4">
|
||||
<span className="font-medium text-sm">Hosts, current minute</span>
|
||||
<Button variant="outline" size="sm" onClick={refresh} disabled={loading}>
|
||||
<RefreshCwIcon data-icon="inline-start" className={loading ? "animate-spin" : undefined} />
|
||||
Refresh
|
||||
</Button>
|
||||
</div>
|
||||
{status?.hosts.length ? (
|
||||
<Table aria-label="Noise counters per host">
|
||||
<TableHeader>
|
||||
<TableRow>
|
||||
<TableHead>Host</TableHead>
|
||||
<TableHead className="text-right">RPM</TableHead>
|
||||
<TableHead className="text-right">Peak RPM</TableHead>
|
||||
<TableHead className="text-right">In flight (peak)</TableHead>
|
||||
<TableHead className="text-right">Duplicates</TableHead>
|
||||
</TableRow>
|
||||
</TableHeader>
|
||||
<TableBody>
|
||||
{status.hosts.map((h) => (
|
||||
<TableRow key={h.host}>
|
||||
<TableCell className="max-w-48 truncate font-mono text-xs" title={h.host}>
|
||||
{h.host}
|
||||
</TableCell>
|
||||
<TableCell className="text-right">{h.rpm}</TableCell>
|
||||
<TableCell className="text-right">{h.peak_rpm}</TableCell>
|
||||
<TableCell className="text-right">
|
||||
{h.inflight} ({h.peak_inflight})
|
||||
</TableCell>
|
||||
<TableCell className="text-right">{h.duplicates_avoided}</TableCell>
|
||||
</TableRow>
|
||||
))}
|
||||
</TableBody>
|
||||
</Table>
|
||||
) : (
|
||||
<p className="text-muted-foreground text-sm">
|
||||
{status?.profile === "off" ? "Off counts no request." : "No paced request yet."}
|
||||
</p>
|
||||
)}
|
||||
</CardContent>
|
||||
</Card>
|
||||
);
|
||||
}
|
||||
@@ -15,6 +15,8 @@ import { api } from "@/lib/api";
|
||||
import { CHAT_SEND_MODE_OPTIONS, type ChatSendMode, setChatSendMode, useChatSendMode } from "@/lib/chat-send-mode";
|
||||
import type { Settings } from "@/lib/types";
|
||||
|
||||
import { NoiseProfileCard } from "./noise-profile-card";
|
||||
|
||||
export default function SystemSettingsPage() {
|
||||
const [trafficCapture, setTrafficCapture] = React.useState(false);
|
||||
const [agentTrafficBinding, setAgentTrafficBinding] = React.useState(false);
|
||||
@@ -292,6 +294,8 @@ export default function SystemSettingsPage() {
|
||||
</CardContent>
|
||||
</Card>
|
||||
|
||||
<NoiseProfileCard />
|
||||
|
||||
<Card className="mb-4 break-inside-avoid md:mb-6">
|
||||
<CardHeader>
|
||||
<CardTitle className="flex items-center gap-2 text-base">
|
||||
|
||||
@@ -826,6 +826,7 @@ export const api = {
|
||||
// ---- app settings (runtime toggles) ----
|
||||
settings: () => get<Settings>(`/settings`),
|
||||
setSettings: (patch: Partial<Settings>) => put<Settings>(`/settings`, patch),
|
||||
noiseStatus: () => get<import("@/lib/types").NoiseStatus>(`/settings/noise`),
|
||||
// Run a real "test" search with the given (or saved) config to verify it works.
|
||||
testWebSearch: (patch: {
|
||||
web_search_backend?: string;
|
||||
|
||||
@@ -962,6 +962,23 @@ export interface Settings {
|
||||
// Render limits of IM messages; 0 = no truncation. The backend answers 400 outside the bounds.
|
||||
notify_summary_max_runes?: number; // summary characters per finding, 0-4000, default 120
|
||||
notify_assets_max_shown?: number; // asset names listed per finding, 0-50, default 3
|
||||
// Target noise profile; the stored setting wins over ARTEX_NOISE_PROFILE. Applies to the next request.
|
||||
noise_profile?: NoiseProfileName;
|
||||
}
|
||||
|
||||
export type NoiseProfileName = "off" | "quiet" | "balanced" | "aggressive";
|
||||
|
||||
// GET /api/settings/noise: the current profile and its per-host counters (current minute, no reset).
|
||||
export interface NoiseStatus {
|
||||
profile: NoiseProfileName;
|
||||
hosts: {
|
||||
host: string;
|
||||
rpm: number;
|
||||
peak_rpm: number;
|
||||
peak_inflight: number;
|
||||
inflight: number;
|
||||
duplicates_avoided: number;
|
||||
}[];
|
||||
}
|
||||
|
||||
// ---- Vulnerability IM push ----
|
||||
|
||||
Reference in New Issue
Block a user