mirror of
https://github.com/BaaRai/ARTEX.git
synced 2026-11-15 04:30:55 +08:00
Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
2853f085e5 |
@@ -1,23 +0,0 @@
|
||||
package agent
|
||||
|
||||
import "strings"
|
||||
import "testing"
|
||||
|
||||
// The deadline contract must stay in the executor seeds: no work whose
|
||||
// timeout can outlive the remaining budget; near the end, land instead of
|
||||
// open. Pattern adapted from bug-bounty-hunting-prompts (MIT).
|
||||
func TestPromptSeedsCarryDeadlineContract(t *testing.T) {
|
||||
for name, tmpl := range BuiltinPromptSeeds() {
|
||||
switch name {
|
||||
case "worker", "pentest":
|
||||
for _, marker := range []string{
|
||||
"timeout can outlive",
|
||||
"prefer the smaller variant that lands",
|
||||
} {
|
||||
if !strings.Contains(tmpl, marker) {
|
||||
t.Errorf("%s seed missing deadline-contract marker %q", name, marker)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -68,7 +68,6 @@ Traffic binding is optional: for TCP and other non-HTTP vulnerabilities, or when
|
||||
|
||||
━━ Judgment and wrap-up ━━
|
||||
- Check against the task goal at any time: when results you have **verified** satisfy the goal, judge it met on that basis and state the grounds. The prerequisite for calling it "met" is that doctrine 4's self-check has passed — an unverified trophy doesn't count as grounds.
|
||||
- **Deadline contract**: before launching any long command (broad scans, wordlist runs, big fetches), check it can finish within the task's remaining time; never start work whose timeout can outlive the task deadline — prefer the smaller variant that lands. As the deadline approaches, shift from discovery to consolidation: stop opening routes, verify and land what exists.
|
||||
- **Wrap-up takes top priority**: when you receive the wrap-up signal (or judge on your own that the goal is met / all reasonable routes are exhausted), **immediately stop all probing and commands**, land the conclusions in hand, and give a concise summary — at this point every prior instruction ("keep exploring / try once more / exhaust this chain / wait for command results") is overridden by wrap-up; do not start any new action.
|
||||
- Summarize in plain language: what was achieved, which routes were taken, which vulnerabilities were confirmed (with PoC locations), which directions are blocked and why. Report only what was really done; never fabricate.
|
||||
|
||||
|
||||
@@ -243,8 +243,6 @@ const workerDefaultTmpl = `You are an "executor" (work agent) of a network-secur
|
||||
- **Confirmed vulnerability → report_finding (exploration graph, with PoC, pass intent_id)**: **use it only when you actually triggered it this time and obtained reproducible evidence (request/response or command output)**. It is strictly forbidden to treat "version/fingerprint matching a CVE", "the parameter looks injectable", or "inference from external vulnerability databases/changelogs/code diffs" as confirmed, and don't substitute querying a CVE database or comparing patch versions for actually triggering it. Can't trigger it but suspicion remains → record an inferred fact with record_fact (the suspicion point + why it couldn't be triggered) for the planner; don't force it into a finding.
|
||||
|
||||
|
||||
**Budget discipline (the deadline contract)**: your run has a step budget and a wall-clock budget. Before starting any command or probe, check it can finish well within what remains; never start work whose timeout can outlive the remaining budget — prefer the smaller variant that lands. When the budget is nearly spent or the wrap-up signal arrives, stop opening new work and land everything already collected (facts first, then findings): a result that exists only in your head is lost.
|
||||
|
||||
After completing this intent, summarize in one sentence what you did and which facts you wrote back.
|
||||
|
||||
Adversarial gate before report_finding:
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
@@ -160,6 +161,10 @@ ORDER BY is_default DESC, priority DESC, id ASC`)
|
||||
|
||||
// SaveProfile inserts (id==0) or updates a profile. Empty apiKey on update keeps existing.
|
||||
func (d *DB) SaveProfile(p *LLMProfile) (int64, error) {
|
||||
// Pasted keys can carry surrounding whitespace; every reader that goes
|
||||
// through ConfigFrom trims, so store the trimmed form or the raw value
|
||||
// leaks into paths that assign it directly (the llm/test fallback).
|
||||
p.APIKey = strings.TrimSpace(p.APIKey)
|
||||
hint := p.APIKeyHint
|
||||
if len(p.APIKey) >= 4 {
|
||||
hint = "…" + p.APIKey[len(p.APIKey)-4:]
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// The llm/test stored-key fallback must send the same header bytes the runtime
|
||||
// sends: form keys are trimmed inside ConfigFrom, the DB fallback used to be
|
||||
// assigned raw, so a key saved with surrounding whitespace ran fine in
|
||||
// conversations but 401'd the connection test.
|
||||
func TestLLMTestStoredKeyFallbackIsTrimmed(t *testing.T) {
|
||||
m, err := NewManager(t.TempDir(), "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { m.Close() })
|
||||
if m.pg == nil {
|
||||
t.Skip("no postgres for this run")
|
||||
}
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
s := New(ctx, m, t.TempDir(), t.TempDir(), t.TempDir())
|
||||
|
||||
var gotKey atomic.Value
|
||||
fake := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotKey.Store(r.Header.Get("x-api-key"))
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"id":"m","type":"message","role":"assistant","model":"m","content":[{"type":"text","text":"OK"}],"stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":1}}`))
|
||||
}))
|
||||
t.Cleanup(fake.Close)
|
||||
|
||||
pid, err := m.pg.SaveProfile(&db.LLMProfile{
|
||||
Name: "ws-key", Format: "anthropic", Model: "claude-x",
|
||||
BaseURL: fake.URL, APIKey: "k1 \t",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _, _ = m.pg.Exec("DELETE FROM llm_profiles WHERE id=$1", pid) })
|
||||
|
||||
// The handler trims the fallback key: the profile was saved through
|
||||
// SaveProfile (which now trims), so also prove the storage trim itself.
|
||||
stored, err := m.pg.ProfileByID(pid)
|
||||
if err != nil || stored == nil {
|
||||
t.Fatalf("ProfileByID: %v %v", stored, err)
|
||||
}
|
||||
if stored.APIKey != "k1" {
|
||||
t.Fatalf("SaveProfile must trim surrounding whitespace, got %q", stored.APIKey)
|
||||
}
|
||||
|
||||
body, _ := json.Marshal(map[string]any{
|
||||
"provider": "anthropic", "model": "claude-x", "base_url": fake.URL,
|
||||
"api_key": "", "profile_id": pid, "streaming": false,
|
||||
})
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/llm/test", bytes.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
rec := httptest.NewRecorder()
|
||||
s.testLLM(rec, req)
|
||||
|
||||
if k, _ := gotKey.Load().(string); k != "k1" {
|
||||
t.Fatalf("test request must send the trimmed stored key, header was %q", k)
|
||||
}
|
||||
// The reply body only proves the call completed; the header is the assertion.
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("handler status %d", rec.Code)
|
||||
}
|
||||
var resp struct {
|
||||
OK bool `json:"ok"`
|
||||
Error string `json:"error"`
|
||||
}
|
||||
_ = json.Unmarshal(rec.Body.Bytes(), &resp)
|
||||
if !resp.OK {
|
||||
t.Logf("note: provider round-trip returned ok=false (%s); header assertion stands", resp.Error)
|
||||
}
|
||||
}
|
||||
|
||||
// Sanity: the runtime path and the test path must agree byte-for-byte on the
|
||||
// key they would send for the same stored profile.
|
||||
func TestLLMTestFallbackMatchesRuntimeNormalization(t *testing.T) {
|
||||
const raw = " k1 \n"
|
||||
runtime := agent.ConfigFrom("anthropic", "m", "https://x.example", raw, "")
|
||||
if runtime.APIKey != "k1" {
|
||||
t.Fatalf("ConfigFrom trims: %q", runtime.APIKey)
|
||||
}
|
||||
}
|
||||
+5
-2
@@ -1404,11 +1404,14 @@ func (s *Server) testLLM(w http.ResponseWriter, r *http.Request) {
|
||||
// Same for the session header name: when the form omits it, the stored profile's value is the fallback.
|
||||
if req.ProfileID != nil && (cfg.APIKey == "" || cfg.SessionHeaderKey == "") {
|
||||
if p, err := s.m.pg.ProfileByID(*req.ProfileID); err == nil && p != nil {
|
||||
// Trim like ConfigFrom does for form values: a key stored with
|
||||
// surrounding whitespace runs fine in conversations (the runtime
|
||||
// path re-trims) but would be sent raw here and fail auth.
|
||||
if cfg.APIKey == "" {
|
||||
cfg.APIKey = p.APIKey
|
||||
cfg.APIKey = strings.TrimSpace(p.APIKey)
|
||||
}
|
||||
if cfg.SessionHeaderKey == "" {
|
||||
cfg.SessionHeaderKey = p.SessionHeaderKey
|
||||
cfg.SessionHeaderKey = strings.TrimSpace(p.SessionHeaderKey)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user