diff --git a/.gitattributes b/.gitattributes index dfdb8b77..eaeb469b 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1 +1,5 @@ *.sh text eol=lf +*.go text eol=lf +*.golden text eol=lf +go.mod text eol=lf +go.sum text eol=lf diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 16baa277..730974d9 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -70,11 +70,38 @@ jobs: exit 1 fi + quality: + runs-on: ubuntu-22.04 + steps: + - name: Checkout + uses: actions/checkout@v6 + with: + fetch-depth: 0 + submodules: recursive + + - name: Set up Go + uses: actions/setup-go@v6 + with: + go-version-file: go.mod + cache: true + + - name: Check dependency layers, registered skips, and repository debt + run: go test -count=1 ./core/deps + + - name: Run go vet + run: go vet ./... + + - name: Check whitespace and submodule pins + run: | + git diff --check HEAD + git diff --exit-code --submodule=diff + git submodule foreach --recursive 'test -z "$(git status --porcelain)"' + # ── Unit tests (depends on tidy) ────────────────────────────── test: runs-on: ubuntu-22.04 - needs: tidy + needs: [tidy, quality] steps: - name: Checkout uses: actions/checkout@v6 @@ -152,13 +179,66 @@ jobs: run: | go test -tags "re2_cgo re2_static" -race -count=1 -timeout 5m -v \ -run 'MultiRound|SendCtrlC' \ - ./pkg/agent/tmux/ + ./agent/tmux/ - name: Run agent tmux integration tests run: | go test -tags "re2_cgo re2_static" -race -count=1 -timeout 5m -v \ -run 'AgentTmux' \ - ./pkg/agent/ + ./agent/ + + windows-test: + runs-on: windows-2022 + needs: [tidy, quality] + steps: + - name: Checkout + uses: actions/checkout@v6 + with: + fetch-depth: 0 + submodules: recursive + + - name: Set up Go + uses: actions/setup-go@v6 + with: + go-version-file: go.mod + cache: true + + - name: Run native Windows package tests + run: go test -count=1 ./agent/... ./pkg/runner/... ./pkg/web/... + + - name: Compile and test the full CLI on Windows + run: go test -count=1 -tags full ./cmd/aiscan + + race-stress: + runs-on: ubuntu-22.04 + needs: [tidy, quality] + steps: + - name: Checkout + uses: actions/checkout@v6 + with: + fetch-depth: 0 + submodules: recursive + + - name: Set up Go + uses: actions/setup-go@v6 + with: + go-version-file: go.mod + cache: true + + - name: Repeat agent and runner concurrency tests + run: | + go test -race -count=20 -timeout 15m \ + -run 'Test(ConcurrentEmitWhileRegistering|SetProviderRaceWithRun|ResetDoesNotAllowConcurrentPrompt|StreamingProviderEmitsMessageUpdates)$' \ + ./agent/... + go test -race -count=20 -timeout 15m \ + -run 'Test(StdioSameSessionFIFOOrder|StdioSessionsRunConcurrently|StdioDrainWaitsForInFlightAndQueued|RuntimeSessionDirectLoopUsesSessionScheduler|RuntimeSessionRejectsRequestsPastPendingLimit|SessionContextCancellationStopsActiveRun|ActiveRunSteersAsyncInputWithoutSecondLifecycle)$' \ + ./pkg/runner/... + + - name: Repeat web SSE, cancellation, and reload concurrency tests + run: | + go test -race -count=20 -timeout 20m \ + -run 'Test(BroadcastAOPEventPersistsRawEnvelope|ServeSSEWithSnapshotSubscribesBeforeReadingSnapshot|ServeSSEWithSnapshotDropsQueuedSnapshotDuplicates|SessionEventsReplayHasNoSideEffects|SessionEventsResumesAfterLastEventID|CancelRemoteScanStopsAgentAndPreservesCanceledStatus|CancelQueuedScanDoesNotWaitForConcurrencySlot|CancelTaskUsesControlChannelWhenTaskQueueIsFull|CancelTaskWaitsForSaturatedControlChannel|CompleteJobCannotOverwriteCanceledScan|BroadcastConfigReload|BroadcastConfigReloadWaitsBehindCancellationFrames|HandleConfigReloadResultUpdatesAgentStatus|SaveConfigBuildFailureKeepsCommittedConfigAndCurrentApp|SaveConfigCommitFailureClosesCandidateAndKeepsCurrentApp|SaveConfigSerializesConcurrentCandidates)$' \ + ./pkg/web scanner-functional: runs-on: ubuntu-22.04 @@ -209,6 +289,9 @@ jobs: go test -tags "re2_cgo re2_static" -race -count=1 -timeout 5m \ ./core/resources/... + - name: Check generated resources are committed + run: git diff --exit-code + # ── E2E tests (depends on test) ─────────────────────────────── e2e: @@ -240,6 +323,20 @@ jobs: npm --prefix web/frontend run build test -s web/static/index.html + - name: Install Playwright Chromium + working-directory: web/frontend + run: npx playwright install --with-deps chromium + + - name: Run frontend Playwright E2E + working-directory: web/frontend + run: npm run test:e2e + + - name: Run cyber-ui viewer tests + working-directory: web/frontend/cyber-ui + run: | + corepack pnpm install --frozen-lockfile + corepack pnpm --filter @cyber/viewer test + - name: Run e2e tests run: | go test -race -count=1 -timeout 10m \ @@ -303,7 +400,7 @@ jobs: - name: Build all platforms (${{ matrix.id }}) run: | - for target in linux/amd64 linux/arm64 darwin/amd64 darwin/arm64 windows/amd64; do + for target in linux/amd64 linux/arm64 darwin/amd64 darwin/arm64 windows/amd64 windows/arm64; do IFS='/' read -r goos goarch <<< "$target" echo " compile ${goos}/${goarch}" CGO_ENABLED=0 GOOS="$goos" GOARCH="$goarch" \ diff --git a/.github/workflows/scanner-regression.yml b/.github/workflows/scanner-regression.yml index 883d80f2..b7c255b0 100644 --- a/.github/workflows/scanner-regression.yml +++ b/.github/workflows/scanner-regression.yml @@ -35,4 +35,4 @@ jobs: run: | go test -tags "full integration re2_cgo re2_static" -count=1 -timeout 8m -v \ -run 'Test(ScannerPublicIntegration|FullScannerPublicIntegration)$' \ - ./pkg/tools + ./tools diff --git a/.gitignore b/.gitignore index cb364dd4..1314135c 100644 --- a/.gitignore +++ b/.gitignore @@ -27,6 +27,8 @@ out/ scan_results.jsonl pw_driver_bin node_modules/ +web/frontend/playwright-report/ +web/frontend/test-results/ community.yaml # Local runtime state / operator artifacts diff --git a/Makefile b/Makefile index 3691d9f3..b22aead0 100644 --- a/Makefile +++ b/Makefile @@ -57,7 +57,7 @@ prepare: mkdir -p "$(BIN_DIR)" aop-gen: - $(GO) generate ./pkg/aop/... + $(GO) generate ./core/aop/... frontend: $(NPM) --prefix "$(WEB_DIR)" run build diff --git a/pkg/agent/agent.go b/agent/agent.go similarity index 95% rename from pkg/agent/agent.go rename to agent/agent.go index 388a9105..f14b92bd 100644 --- a/pkg/agent/agent.go +++ b/agent/agent.go @@ -5,9 +5,9 @@ import ( "fmt" "sync" - "github.com/chainreactors/aiscan/pkg/agent/inbox" - "github.com/chainreactors/aiscan/pkg/aop/x/delegation" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/agent/inbox" + "github.com/chainreactors/aiscan/core/aop/x/delegation" + "github.com/chainreactors/aiscan/core/telemetry" ) type Agent struct { @@ -81,17 +81,15 @@ func (a *Agent) SessionID() string { } func (a *Agent) beginSession() { - a.mu.Lock() - em, model := a.Cfg.emitter, a.Cfg.Model - a.mu.Unlock() - em.sessionStart(model) + cfg := a.configSnapshot() + cfg.emitter.sessionStart(cfg.Model) + emitSessionStart(context.Background(), cfg) } func (a *Agent) endSession(reason string) { - a.mu.Lock() - em := a.Cfg.emitter - a.mu.Unlock() - em.sessionEnd(reason) + cfg := a.configSnapshot() + cfg.emitter.sessionEnd(reason) + emitSessionEnd(context.Background(), cfg, reason) } // Continue resumes the agent without a new prompt (e.g. after tool results). @@ -235,6 +233,7 @@ func deriveNamedFromConfig(cfg Config, name, parentToolCallID string, detail *de Temperature: cfg.Temperature, CacheRetention: cfg.CacheRetention, Bus: cfg.Bus, + Hooks: cfg.Hooks, AgentName: name, ParentSessionID: cfg.SessionID, ParentToolCallID: parentToolCallID, diff --git a/pkg/agent/agent_test.go b/agent/agent_test.go similarity index 99% rename from pkg/agent/agent_test.go rename to agent/agent_test.go index 5af484c8..637c77de 100644 --- a/pkg/agent/agent_test.go +++ b/agent/agent_test.go @@ -12,9 +12,9 @@ import ( "testing" "time" - "github.com/chainreactors/aiscan/pkg/aop" + "github.com/chainreactors/aiscan/core/aop" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" "github.com/chainreactors/aiscan/skills" ) diff --git a/pkg/agent/aop_emit.go b/agent/aop_emit.go similarity index 98% rename from pkg/agent/aop_emit.go rename to agent/aop_emit.go index 7149fc9a..173717f8 100644 --- a/pkg/agent/aop_emit.go +++ b/agent/aop_emit.go @@ -6,9 +6,9 @@ import ( "sync/atomic" "time" + "github.com/chainreactors/aiscan/core/aop" + "github.com/chainreactors/aiscan/core/aop/x/delegation" "github.com/chainreactors/aiscan/core/eventbus" - "github.com/chainreactors/aiscan/pkg/aop" - "github.com/chainreactors/aiscan/pkg/aop/x/delegation" ) // aopEmitter is the agent kernel's single event-emission path. Every event diff --git a/pkg/agent/aop_emit_test.go b/agent/aop_emit_test.go similarity index 99% rename from pkg/agent/aop_emit_test.go rename to agent/aop_emit_test.go index 9de95842..d858286a 100644 --- a/pkg/agent/aop_emit_test.go +++ b/agent/aop_emit_test.go @@ -6,7 +6,7 @@ import ( "sync/atomic" "testing" - "github.com/chainreactors/aiscan/pkg/aop" + "github.com/chainreactors/aiscan/core/aop" ) // streamEventCollector records message/message.delta events from the bus. diff --git a/pkg/agent/compact.go b/agent/compact.go similarity index 98% rename from pkg/agent/compact.go rename to agent/compact.go index bf79ff40..e1b6786b 100644 --- a/pkg/agent/compact.go +++ b/agent/compact.go @@ -5,8 +5,8 @@ import ( "fmt" "strings" - "github.com/chainreactors/aiscan/pkg/agent/truncate" - xcompact "github.com/chainreactors/aiscan/pkg/aop/x/compact" + xcompact "github.com/chainreactors/aiscan/core/aop/x/compact" + "github.com/chainreactors/aiscan/core/truncate" ) const compactSystemPrompt = `You are a context summarization assistant. Your task is to read a conversation between a user and an AI assistant, then produce a structured summary following the exact format specified. diff --git a/pkg/agent/compact_test.go b/agent/compact_test.go similarity index 100% rename from pkg/agent/compact_test.go rename to agent/compact_test.go diff --git a/pkg/agent/context_window.go b/agent/context_window.go similarity index 100% rename from pkg/agent/context_window.go rename to agent/context_window.go diff --git a/pkg/agent/cron.go b/agent/cron.go similarity index 100% rename from pkg/agent/cron.go rename to agent/cron.go diff --git a/pkg/agent/cron_test.go b/agent/cron_test.go similarity index 100% rename from pkg/agent/cron_test.go rename to agent/cron_test.go diff --git a/pkg/agent/defaults.go b/agent/defaults.go similarity index 87% rename from pkg/agent/defaults.go rename to agent/defaults.go index b3d55c7a..1b6a392e 100644 --- a/pkg/agent/defaults.go +++ b/agent/defaults.go @@ -1,6 +1,6 @@ package agent -import "github.com/chainreactors/aiscan/pkg/agent/truncate" +import "github.com/chainreactors/aiscan/core/truncate" const ( DefaultMaxResultSize = truncate.DefaultMaxBytes diff --git a/pkg/agent/evaluator/evaluator.go b/agent/evaluator/evaluator.go similarity index 96% rename from pkg/agent/evaluator/evaluator.go rename to agent/evaluator/evaluator.go index 734ab26f..e4608ac5 100644 --- a/pkg/agent/evaluator/evaluator.go +++ b/agent/evaluator/evaluator.go @@ -7,10 +7,10 @@ import ( "strings" "time" - agentpkg "github.com/chainreactors/aiscan/pkg/agent" - "github.com/chainreactors/aiscan/pkg/agent/provider" - "github.com/chainreactors/aiscan/pkg/agent/truncate" - "github.com/chainreactors/aiscan/pkg/telemetry" + agentpkg "github.com/chainreactors/aiscan/agent" + "github.com/chainreactors/aiscan/agent/provider" + "github.com/chainreactors/aiscan/core/telemetry" + "github.com/chainreactors/aiscan/core/truncate" ) const ( diff --git a/pkg/agent/evaluator/loop.go b/agent/evaluator/loop.go similarity index 96% rename from pkg/agent/evaluator/loop.go rename to agent/evaluator/loop.go index 0558bc34..c41febb9 100644 --- a/pkg/agent/evaluator/loop.go +++ b/agent/evaluator/loop.go @@ -5,10 +5,10 @@ import ( "fmt" "strings" - "github.com/chainreactors/aiscan/pkg/agent" - "github.com/chainreactors/aiscan/pkg/agent/provider" - xeval "github.com/chainreactors/aiscan/pkg/aop/x/eval" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/agent" + "github.com/chainreactors/aiscan/agent/provider" + xeval "github.com/chainreactors/aiscan/core/aop/x/eval" + "github.com/chainreactors/aiscan/core/telemetry" ) const defaultMaxEvalRounds = 3 diff --git a/pkg/agent/evaluator/loop_test.go b/agent/evaluator/loop_test.go similarity index 95% rename from pkg/agent/evaluator/loop_test.go rename to agent/evaluator/loop_test.go index 9877cbe2..cd63ac0e 100644 --- a/pkg/agent/evaluator/loop_test.go +++ b/agent/evaluator/loop_test.go @@ -4,10 +4,10 @@ import ( "context" "testing" + "github.com/chainreactors/aiscan/agent" + "github.com/chainreactors/aiscan/agent/provider" + "github.com/chainreactors/aiscan/core/aop" "github.com/chainreactors/aiscan/core/eventbus" - "github.com/chainreactors/aiscan/pkg/agent" - "github.com/chainreactors/aiscan/pkg/agent/provider" - "github.com/chainreactors/aiscan/pkg/aop" ) type fixedProvider struct { diff --git a/agent/external_import_test.go b/agent/external_import_test.go new file mode 100644 index 00000000..d4861f14 --- /dev/null +++ b/agent/external_import_test.go @@ -0,0 +1,18 @@ +package agent_test + +import ( + "testing" + + "github.com/chainreactors/aiscan/agent" +) + +func TestRootAgentPublicImport(t *testing.T) { + config := agent.Config{}. + WithModel("example-model"). + WithMaxTokens(256). + WithContextWindow(4096) + if config.Model != "example-model" || config.MaxTokens != 256 || config.ContextWindow != 4096 { + t.Fatalf("root agent config aliases/builders are not externally usable: %#v", config) + } + _ = agent.ProviderConfig{Model: "example-model"} +} diff --git a/pkg/agent/finish_tool.go b/agent/finish_tool.go similarity index 100% rename from pkg/agent/finish_tool.go rename to agent/finish_tool.go diff --git a/pkg/agent/helpers_test.go b/agent/helpers_test.go similarity index 99% rename from pkg/agent/helpers_test.go rename to agent/helpers_test.go index 6416f91a..e2850e77 100644 --- a/pkg/agent/helpers_test.go +++ b/agent/helpers_test.go @@ -10,10 +10,10 @@ import ( "sync/atomic" "testing" + "github.com/chainreactors/aiscan/agent/inbox" + "github.com/chainreactors/aiscan/core/aop" "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/tool" - "github.com/chainreactors/aiscan/pkg/agent/inbox" - "github.com/chainreactors/aiscan/pkg/aop" "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/aiscan/skills" ) diff --git a/agent/hooks/hooks.go b/agent/hooks/hooks.go new file mode 100644 index 00000000..816fc143 --- /dev/null +++ b/agent/hooks/hooks.go @@ -0,0 +1,310 @@ +// Package hooks is the agent kernel's single extension mechanism: typed hook +// points with explicit result semantics and error policies. +// +// A hook point is a package-level Point[E, R] descriptor carrying both type +// parameters, so callers write ToolCallHook.Emit(ctx, reg, ev) without spelling +// out E and R. The Registry stores handlers type-erased and is copy-on-write: +// registration takes a mutex, dispatch is a single atomic load. Tool calls run +// from N goroutines concurrently, so Emit must never block on registration. +package hooks + +import ( + "context" + "errors" + "fmt" + "sync" + "sync/atomic" +) + +type Kind string + +type ErrorPolicy uint8 + +const ( + ContinueOnError ErrorPolicy = iota // collect, report, keep dispatching + FailClosed // first error aborts dispatch; caller must deny +) + +// HandlerError attributes a failure to the handler that produced it. Handlers +// are registered with a mandatory source so a hook failure never has to be +// traced back by hand. +type HandlerError struct { + Source string + Kind Kind + Err error +} + +func (e *HandlerError) Error() string { + return fmt.Sprintf("hook %s/%s: %v", e.Kind, e.Source, e.Err) +} + +func (e *HandlerError) Unwrap() error { return e.Err } + +// errTypeMismatch means two Points share a Kind with different E/R. Reporting it +// as a handler failure keeps the handler from silently vanishing. +var errTypeMismatch = errors.New("handler signature does not match hook point") + +// Reducer folds one handler result into the accumulated result. ev is a pointer +// so fold-style points can let the next handler observe the previous handler's +// change. Returning true short-circuits the remaining handlers. +type Reducer[E any, R any] func(acc *R, ev *E, out R) (stop bool) + +type Point[E any, R any] struct { + Kind Kind + Reduce Reducer[E, R] // nil => pure observation + OnError ErrorPolicy +} + +type entry struct { + id uint64 + source string + fn any // func(context.Context, E) (R, error), asserted back in dispatch +} + +// table is replaced wholesale on every registration change; readers only ever +// see a consistent immutable snapshot. +type table struct { + byKind map[Kind][]entry +} + +type errorSink struct { + fn func(*HandlerError) +} + +type cleanup struct { + id uint64 + fn func() +} + +type Registry struct { + mu sync.Mutex + nextID uint64 + cleanups []cleanup + + handlers atomic.Pointer[table] + sink atomic.Pointer[errorSink] +} + +func New() *Registry { + r := &Registry{} + r.handlers.Store(&table{byKind: map[Kind][]entry{}}) + return r +} + +// Has is the zero-handler fast path: one atomic load plus a map lookup, no locks +// and no allocations. +func (r *Registry) Has(kind Kind) bool { + return r.Len(kind) > 0 +} + +func (r *Registry) Len(kind Kind) int { + if r == nil { + return 0 + } + t := r.handlers.Load() + if t == nil { + return 0 + } + return len(t.byKind[kind]) +} + +// SetErrorSink installs the reporter for handler failures. Passing nil disables +// reporting; errors are still collected and returned by Emit. +func (r *Registry) SetErrorSink(fn func(*HandlerError)) { + if r == nil { + return + } + r.sink.Store(&errorSink{fn: fn}) +} + +func (r *Registry) report(he *HandlerError) { + s := r.sink.Load() + if s == nil || s.fn == nil { + return + } + s.fn(he) +} + +// AddCleanup registers a function to run on Clear. The returned remove is +// idempotent. +func (r *Registry) AddCleanup(fn func()) (remove func()) { + if r == nil || fn == nil { + return func() {} + } + r.mu.Lock() + r.nextID++ + id := r.nextID + r.cleanups = append(r.cleanups, cleanup{id: id, fn: fn}) + r.mu.Unlock() + + var once sync.Once + return func() { + once.Do(func() { + r.mu.Lock() + for i := range r.cleanups { + if r.cleanups[i].id == id { + r.cleanups = append(r.cleanups[:i], r.cleanups[i+1:]...) + break + } + } + r.mu.Unlock() + }) + } +} + +// Clear drops every handler and runs the registered cleanups in registration +// order. +func (r *Registry) Clear() { + if r == nil { + return + } + r.mu.Lock() + r.handlers.Store(&table{byKind: map[Kind][]entry{}}) + pending := r.cleanups + r.cleanups = nil + r.mu.Unlock() + + // Outside the lock so a cleanup may touch the registry itself. + for _, c := range pending { + c.fn() + } +} + +func (r *Registry) add(kind Kind, source string, fn any) func() { + r.mu.Lock() + r.nextID++ + id := r.nextID + next := r.cloneLocked() + prev := next.byKind[kind] + list := make([]entry, len(prev), len(prev)+1) + copy(list, prev) + next.byKind[kind] = append(list, entry{id: id, source: source, fn: fn}) + r.handlers.Store(next) + r.mu.Unlock() + + var once sync.Once + return func() { once.Do(func() { r.remove(kind, id) }) } +} + +func (r *Registry) remove(kind Kind, id uint64) { + r.mu.Lock() + defer r.mu.Unlock() + + old := r.handlers.Load() + if old == nil { + return + } + prev := old.byKind[kind] + idx := -1 + for i := range prev { + if prev[i].id == id { + idx = i + break + } + } + if idx < 0 { + return + } + next := r.cloneLocked() + if len(prev) == 1 { + delete(next.byKind, kind) + } else { + list := make([]entry, 0, len(prev)-1) + list = append(list, prev[:idx]...) + list = append(list, prev[idx+1:]...) + next.byKind[kind] = list + } + r.handlers.Store(next) +} + +// cloneLocked copies the kind map; the per-kind slices stay shared because an +// in-flight dispatch may still be iterating them. +func (r *Registry) cloneLocked() *table { + old := r.handlers.Load() + if old == nil { + return &table{byKind: make(map[Kind][]entry, 1)} + } + next := &table{byKind: make(map[Kind][]entry, len(old.byKind)+1)} + for k, v := range old.byKind { + next.byKind[k] = v + } + return next +} + +// On registers fn for this point and returns an idempotent unsubscribe that is +// safe to call from inside a dispatch. source is mandatory: an unattributable +// handler cannot be reported when it fails, so an empty source (or a nil fn) is +// a programming error and panics. A nil registry is not — hooks are optional +// wiring, so registration on one is a no-op. +func (p Point[E, R]) On(r *Registry, source string, fn func(context.Context, E) (R, error)) (unsubscribe func()) { + if source == "" { + panic("hooks: On requires a non-empty source for " + string(p.Kind)) + } + if fn == nil { + panic("hooks: On requires a non-nil handler for " + string(p.Kind)) + } + if r == nil { + return func() {} + } + return r.add(p.Kind, source, fn) +} + +// Emit runs the point's handlers sequentially in registration order and folds +// their results through Reduce. With no handlers it returns the zero result +// without allocating. +func (p Point[E, R]) Emit(ctx context.Context, r *Registry, ev E) (R, error) { + var zero R + if r == nil { + return zero, nil + } + t := r.handlers.Load() + if t == nil { + return zero, nil + } + entries := t.byKind[p.Kind] + if len(entries) == 0 { + return zero, nil + } + return p.dispatch(ctx, r, entries, ev) +} + +// dispatch is kept out of Emit so that taking &ev here does not force Emit's +// argument onto the heap on the zero-handler path. +// +//go:noinline +func (p Point[E, R]) dispatch(ctx context.Context, r *Registry, entries []entry, ev E) (R, error) { + var acc R + var errs []error + + // entries is an immutable snapshot: handlers may subscribe or unsubscribe + // during dispatch, and this round still sees the set it started with. + for _, e := range entries { + fn, ok := e.fn.(func(context.Context, E) (R, error)) + if !ok { + he := &HandlerError{Source: e.source, Kind: p.Kind, Err: errTypeMismatch} + r.report(he) + errs = append(errs, he) + if p.OnError == FailClosed { + return acc, errors.Join(errs...) + } + continue + } + out, err := fn(ctx, ev) + if err != nil { + he := &HandlerError{Source: e.source, Kind: p.Kind, Err: err} + r.report(he) + errs = append(errs, he) + if p.OnError == FailClosed { + return acc, errors.Join(errs...) + } + continue + } + if p.Reduce == nil { + continue + } + if p.Reduce(&acc, &ev, out) { + break + } + } + return acc, errors.Join(errs...) +} diff --git a/agent/hooks/hooks_test.go b/agent/hooks/hooks_test.go new file mode 100644 index 00000000..e00732ad --- /dev/null +++ b/agent/hooks/hooks_test.go @@ -0,0 +1,458 @@ +package hooks + +import ( + "context" + "errors" + "strings" + "sync" + "sync/atomic" + "testing" +) + +func ptr[T any](v T) *T { return &v } + +func TestEmitRunsHandlersInRegistrationOrder(t *testing.T) { + r := New() + var order []string + for _, name := range []string{"a", "b", "c"} { + Context.On(r, name, func(_ context.Context, _ ContextEvent) (ContextResult, error) { + order = append(order, name) + return ContextResult{}, nil + }) + } + + if _, err := Context.Emit(context.Background(), r, ContextEvent{}); err != nil { + t.Fatalf("emit: %v", err) + } + if got := strings.Join(order, ""); got != "abc" { + t.Fatalf("order = %q, want %q", got, "abc") + } + if n := r.Len("context"); n != 3 { + t.Fatalf("Len = %d, want 3", n) + } +} + +func TestUnsubscribeIsIdempotent(t *testing.T) { + r := New() + var calls int + off := RunEnd.On(r, "counter", func(_ context.Context, _ RunEndEvent) (struct{}, error) { + calls++ + return struct{}{}, nil + }) + + if _, err := RunEnd.Emit(context.Background(), r, RunEndEvent{}); err != nil { + t.Fatalf("emit: %v", err) + } + off() + off() + if r.Has("run_end") { + t.Fatal("Has after unsubscribe = true") + } + if _, err := RunEnd.Emit(context.Background(), r, RunEndEvent{}); err != nil { + t.Fatalf("emit: %v", err) + } + if calls != 1 { + t.Fatalf("calls = %d, want 1", calls) + } +} + +// The in-flight dispatch works off an immutable snapshot, so a handler that +// unsubscribes a later handler does not affect the round already running. +func TestUnsubscribeDuringDispatch(t *testing.T) { + r := New() + var seen []string + var offSecond func() + + Context.On(r, "first", func(_ context.Context, _ ContextEvent) (ContextResult, error) { + seen = append(seen, "first") + offSecond() + return ContextResult{}, nil + }) + offSecond = Context.On(r, "second", func(_ context.Context, _ ContextEvent) (ContextResult, error) { + seen = append(seen, "second") + return ContextResult{}, nil + }) + + if _, err := Context.Emit(context.Background(), r, ContextEvent{}); err != nil { + t.Fatalf("emit: %v", err) + } + if got := strings.Join(seen, ","); got != "first,second" { + t.Fatalf("first dispatch = %q, want %q", got, "first,second") + } + + seen = nil + if _, err := Context.Emit(context.Background(), r, ContextEvent{}); err != nil { + t.Fatalf("emit: %v", err) + } + if got := strings.Join(seen, ","); got != "first" { + t.Fatalf("second dispatch = %q, want %q", got, "first") + } +} + +func TestFailClosedShortCircuits(t *testing.T) { + r := New() + var sunk []*HandlerError + r.SetErrorSink(func(he *HandlerError) { sunk = append(sunk, he) }) + + boom := errors.New("boom") + var secondRan bool + ToolCallHook.On(r, "proxy", func(_ context.Context, _ ToolCallEvent) (ToolCallResult, error) { + return ToolCallResult{}, boom + }) + ToolCallHook.On(r, "audit", func(_ context.Context, _ ToolCallEvent) (ToolCallResult, error) { + secondRan = true + return ToolCallResult{Block: true, Reason: "nope"}, nil + }) + + res, err := ToolCallHook.Emit(context.Background(), r, ToolCallEvent{}) + if err == nil { + t.Fatal("err = nil, want failure") + } + if secondRan { + t.Fatal("second handler ran after fail-closed abort") + } + if res.Block { + t.Fatal("result should be zero when dispatch aborts") + } + if !errors.Is(err, boom) { + t.Fatalf("errors.Is(err, boom) = false: %v", err) + } + + var he *HandlerError + if !errors.As(err, &he) { + t.Fatalf("errors.As(*HandlerError) = false: %v", err) + } + if he.Source != "proxy" || he.Kind != "tool_call" { + t.Fatalf("attribution = %s/%s, want tool_call/proxy", he.Kind, he.Source) + } + if got := he.Error(); got != "hook tool_call/proxy: boom" { + t.Fatalf("Error() = %q", got) + } + if len(sunk) != 1 || sunk[0] != he { + t.Fatalf("sink got %d errors, want the one reported", len(sunk)) + } +} + +func TestContinueOnErrorCollectsAndKeepsGoing(t *testing.T) { + r := New() + first := errors.New("first") + second := errors.New("second") + var ran int + + for _, tc := range []struct { + source string + err error + }{{"a", first}, {"b", second}, {"c", nil}} { + RunEnd.On(r, tc.source, func(_ context.Context, _ RunEndEvent) (struct{}, error) { + ran++ + return struct{}{}, tc.err + }) + } + + _, err := RunEnd.Emit(context.Background(), r, RunEndEvent{}) + if ran != 3 { + t.Fatalf("ran = %d, want 3", ran) + } + if !errors.Is(err, first) || !errors.Is(err, second) { + t.Fatalf("err = %v, want both collected", err) + } +} + +func TestToolResultPatchChaining(t *testing.T) { + r := New() + var observed string + + ToolResult.On(r, "redact", func(_ context.Context, ev ToolResultEvent) (ToolResultPatch, error) { + return ToolResultPatch{Content: ptr(ev.Content + "+redacted")}, nil + }) + ToolResult.On(r, "truncate", func(_ context.Context, ev ToolResultEvent) (ToolResultPatch, error) { + observed = ev.Content + return ToolResultPatch{IsError: ptr(true), Terminate: ptr(true)}, nil + }) + + patch, err := ToolResult.Emit(context.Background(), r, ToolResultEvent{Content: "raw"}) + if err != nil { + t.Fatalf("emit: %v", err) + } + if observed != "raw+redacted" { + t.Fatalf("second handler saw %q, want the first handler's patch", observed) + } + if patch.Content == nil || *patch.Content != "raw+redacted" { + t.Fatalf("patch.Content = %v", patch.Content) + } + if patch.IsError == nil || !*patch.IsError { + t.Fatalf("patch.IsError = %v", patch.IsError) + } + if patch.Terminate == nil || !*patch.Terminate { + t.Fatalf("patch.Terminate = %v", patch.Terminate) + } +} + +func TestBeforeRunFoldsSystemPromptAndAggregatesPrepend(t *testing.T) { + r := New() + var observed string + + BeforeRun.On(r, "base", func(_ context.Context, ev RunStartEvent) (RunStartResult, error) { + return RunStartResult{ + SystemPrompt: ptr(ev.SystemPrompt + "\nbase"), + Prepend: []Msg{{Role: "system", Content: ptr("one")}}, + }, nil + }) + BeforeRun.On(r, "extra", func(_ context.Context, ev RunStartEvent) (RunStartResult, error) { + observed = ev.SystemPrompt + return RunStartResult{ + SystemPrompt: ptr(ev.SystemPrompt + "\nextra"), + Prepend: []Msg{{Role: "user", Content: ptr("two")}}, + }, nil + }) + + res, err := BeforeRun.Emit(context.Background(), r, RunStartEvent{SystemPrompt: "root"}) + if err != nil { + t.Fatalf("emit: %v", err) + } + if observed != "root\nbase" { + t.Fatalf("second handler saw %q, want the folded prompt", observed) + } + if res.SystemPrompt == nil || *res.SystemPrompt != "root\nbase\nextra" { + t.Fatalf("SystemPrompt = %v", res.SystemPrompt) + } + if len(res.Prepend) != 2 || res.Prepend[0].Role != "system" || res.Prepend[1].Role != "user" { + t.Fatalf("Prepend = %+v", res.Prepend) + } +} + +func TestContextReplacementFolds(t *testing.T) { + r := New() + var observed int + + Context.On(r, "drop", func(_ context.Context, ev ContextEvent) (ContextResult, error) { + return ContextResult{Messages: ev.Messages[1:]}, nil + }) + Context.On(r, "noop", func(_ context.Context, ev ContextEvent) (ContextResult, error) { + observed = len(ev.Messages) + return ContextResult{}, nil + }) + + res, err := Context.Emit(context.Background(), r, ContextEvent{Messages: make([]Msg, 3)}) + if err != nil { + t.Fatalf("emit: %v", err) + } + if observed != 2 { + t.Fatalf("second handler saw %d messages, want 2", observed) + } + if len(res.Messages) != 2 { + t.Fatalf("result = %d messages, want 2", len(res.Messages)) + } +} + +func TestStopWhenShortCircuits(t *testing.T) { + r := New() + var ran int + + BeforeCompact.On(r, "budget", func(_ context.Context, _ CompactEvent) (CancelResult, error) { + ran++ + return CancelResult{Cancel: true, Reason: "still cheap"}, nil + }) + BeforeCompact.On(r, "never", func(_ context.Context, _ CompactEvent) (CancelResult, error) { + ran++ + return CancelResult{}, nil + }) + + res, err := BeforeCompact.Emit(context.Background(), r, CompactEvent{}) + if err != nil { + t.Fatalf("emit: %v", err) + } + if ran != 1 { + t.Fatalf("ran = %d, want 1", ran) + } + if !res.Cancel || res.Reason != "still cheap" { + t.Fatalf("res = %+v", res) + } +} + +func TestObservationPointsIgnoreResults(t *testing.T) { + r := New() + var ran int + for _, name := range []string{"a", "b"} { + SessionStart.On(r, name, func(_ context.Context, _ SessionEvent) (struct{}, error) { + ran++ + return struct{}{}, nil + }) + } + + res, err := SessionStart.Emit(context.Background(), r, SessionEvent{SessionID: "s1"}) + if err != nil { + t.Fatalf("emit: %v", err) + } + if ran != 2 { + t.Fatalf("ran = %d, want 2 (observation must not short-circuit)", ran) + } + if res != (struct{}{}) { + t.Fatal("observation result must be zero") + } +} + +var ( + sinkResult ToolCallResult + sinkErr error +) + +func TestEmitFastPathDoesNotAllocate(t *testing.T) { + r := New() + // A handler on a different kind ensures the map lookup misses rather than + // short-circuiting on an empty table. + RunEnd.On(r, "other", func(_ context.Context, _ RunEndEvent) (struct{}, error) { + return struct{}{}, nil + }) + if r.Has("tool_call") { + t.Fatal("Has(tool_call) = true") + } + + ctx := context.Background() + ev := ToolCallEvent{SessionID: "s1", TurnID: "t1", Call: ToolCall{ID: "c1"}} + + if got := testing.AllocsPerRun(100, func() { + sinkResult, sinkErr = ToolCallHook.Emit(ctx, r, ev) + }); got != 0 { + t.Fatalf("Emit allocs = %v, want 0", got) + } + if sinkErr != nil || sinkResult.Block { + t.Fatalf("fast path returned %+v, %v", sinkResult, sinkErr) + } + + if got := testing.AllocsPerRun(100, func() { + sinkResult, sinkErr = ToolCallHook.Emit(ctx, nil, ev) + }); got != 0 { + t.Fatalf("nil-registry Emit allocs = %v, want 0", got) + } +} + +func TestNilRegistryTolerated(t *testing.T) { + var r *Registry + if r.Has("tool_call") || r.Len("tool_call") != 0 { + t.Fatal("nil registry reports handlers") + } + r.SetErrorSink(func(*HandlerError) {}) + r.Clear() + off := ToolCallHook.On(r, "x", func(_ context.Context, _ ToolCallEvent) (ToolCallResult, error) { + return ToolCallResult{}, nil + }) + off() + + res, err := ToolCallHook.Emit(context.Background(), r, ToolCallEvent{}) + if err != nil || res.Block { + t.Fatalf("nil registry Emit = %+v, %v", res, err) + } +} + +func TestOnRequiresSource(t *testing.T) { + defer func() { + if recover() == nil { + t.Fatal("On with empty source did not panic") + } + }() + ToolCallHook.On(New(), "", func(_ context.Context, _ ToolCallEvent) (ToolCallResult, error) { + return ToolCallResult{}, nil + }) +} + +func TestClearDropsHandlersAndRunsCleanups(t *testing.T) { + r := New() + RunEnd.On(r, "a", func(_ context.Context, _ RunEndEvent) (struct{}, error) { + return struct{}{}, nil + }) + + var order []string + r.AddCleanup(func() { order = append(order, "first") }) + removeSecond := r.AddCleanup(func() { order = append(order, "second") }) + r.AddCleanup(func() { order = append(order, "third") }) + removeSecond() + removeSecond() + + r.Clear() + if r.Has("run_end") { + t.Fatal("Clear left handlers behind") + } + if got := strings.Join(order, ","); got != "first,third" { + t.Fatalf("cleanups = %q, want %q", got, "first,third") + } + + r.Clear() + if len(order) != 2 { + t.Fatalf("cleanups ran twice: %v", order) + } +} + +// Two points sharing a Kind with different types must surface as an attributed +// error rather than a silently skipped handler. +func TestSignatureMismatchIsReported(t *testing.T) { + r := New() + imposter := Point[SessionEvent, struct{}]{Kind: ToolCallHook.Kind} + imposter.On(r, "imposter", func(_ context.Context, _ SessionEvent) (struct{}, error) { + return struct{}{}, nil + }) + + _, err := ToolCallHook.Emit(context.Background(), r, ToolCallEvent{}) + if !errors.Is(err, errTypeMismatch) { + t.Fatalf("err = %v, want type mismatch", err) + } +} + +func TestConcurrentEmitWhileRegistering(t *testing.T) { + r := New() + var calls atomic.Int64 + r.SetErrorSink(func(*HandlerError) {}) + + ctx := context.Background() + stop := make(chan struct{}) + var emitters, registrars sync.WaitGroup + + for i := 0; i < 8; i++ { + emitters.Add(1) + go func() { + defer emitters.Done() + for { + select { + case <-stop: + return + default: + } + if _, err := ToolResult.Emit(ctx, r, ToolResultEvent{Content: "x"}); err != nil { + t.Errorf("emit: %v", err) + return + } + _, _ = RunEnd.Emit(ctx, r, RunEndEvent{Stop: StopReasonCompleted}) + } + }() + } + + for i := 0; i < 4; i++ { + registrars.Add(1) + go func() { + defer registrars.Done() + for j := 0; j < 200; j++ { + off := ToolResult.On(r, "racer", func(_ context.Context, ev ToolResultEvent) (ToolResultPatch, error) { + calls.Add(1) + return ToolResultPatch{Content: ptr(ev.Content + "!")}, nil + }) + offEnd := RunEnd.On(r, "racer", func(_ context.Context, _ RunEndEvent) (struct{}, error) { + calls.Add(1) + return struct{}{}, nil + }) + off() + offEnd() + } + }() + } + + registrars.Wait() + close(stop) + emitters.Wait() + + if n := r.Len("tool_result"); n != 0 { + t.Fatalf("leftover handlers: %d", n) + } + if calls.Load() == 0 { + t.Fatal("no handler ever ran concurrently with registration") + } +} diff --git a/agent/hooks/points.go b/agent/hooks/points.go new file mode 100644 index 00000000..9d717423 --- /dev/null +++ b/agent/hooks/points.go @@ -0,0 +1,182 @@ +package hooks + +import ( + "github.com/chainreactors/aiscan/agent/provider" + "github.com/chainreactors/aiscan/core/tool" +) + +// Aliases keep event definitions readable without pulling agent in (that +// would be an import cycle). +type ( + Msg = provider.ChatMessage + ToolCall = provider.ToolCall +) + +// StopReason lives here rather than in agent because run_end events carry +// it; agent aliases these back. +type StopReason string + +const ( + StopReasonCompleted StopReason = "completed" + StopReasonTerminated StopReason = "terminated" + StopReasonStopped StopReason = "stopped" + StopReasonBudget StopReason = "budget" + StopReasonError StopReason = "error" + StopReasonCanceled StopReason = "canceled" +) + +// RunStartEvent carries the config context a handler needs in flattened form, +// since the event type cannot reference *agent.Config. +type RunStartEvent struct { + SessionID string + TurnID string + AgentName string + Model string + Turn int + SystemPrompt string + ToolNames []string +} + +// RunStartResult replaces the system prompt (nil = keep) and prepends messages +// to the turn. +type RunStartResult struct { + SystemPrompt *string + Prepend []Msg +} + +var BeforeRun = Point[RunStartEvent, RunStartResult]{ + Kind: "before_run", + Reduce: Fold(func(acc *RunStartResult, ev *RunStartEvent, out RunStartResult) { + if out.SystemPrompt != nil { + // Fold into the event so the next handler edits the new prompt. + ev.SystemPrompt = *out.SystemPrompt + acc.SystemPrompt = out.SystemPrompt + } + acc.Prepend = append(acc.Prepend, out.Prepend...) + }), +} + +type ContextEvent struct { + SessionID string + Turn int + Messages []Msg +} + +// ContextResult replaces the whole message list; nil means unchanged. +type ContextResult struct { + Messages []Msg +} + +var Context = Point[ContextEvent, ContextResult]{ + Kind: "context", + Reduce: Fold(func(acc *ContextResult, ev *ContextEvent, out ContextResult) { + if out.Messages == nil { + return + } + ev.Messages = out.Messages + acc.Messages = out.Messages + }), +} + +type ToolCallEvent struct { + SessionID string + TurnID string + AssistantMessage Msg + Call ToolCall + SystemPrompt string + Messages []Msg +} + +type ToolCallResult struct { + Block bool + Reason string +} + +// ToolCallHook is fail-closed: a handler that errors out cannot be assumed to +// have approved the call, so the caller must treat any error as a denial. +var ToolCallHook = Point[ToolCallEvent, ToolCallResult]{ + Kind: "tool_call", + OnError: FailClosed, + Reduce: StopWhen[ToolCallEvent](func(r ToolCallResult) bool { return r.Block }), +} + +type ToolResultEvent struct { + SessionID string + TurnID string + Call ToolCall + Content string + IsError bool + Terminate bool + DurationMs int + Full *tool.Result +} + +// ToolResultPatch patches individual fields; nil fields are left alone. +type ToolResultPatch struct { + Content *string + IsError *bool + Terminate *bool +} + +var ToolResult = Point[ToolResultEvent, ToolResultPatch]{ + Kind: "tool_result", + Reduce: Fold(func(acc *ToolResultPatch, ev *ToolResultEvent, out ToolResultPatch) { + // Each patch is mirrored onto the event so later handlers see the + // already-patched result rather than the original. + if out.Content != nil { + ev.Content = *out.Content + acc.Content = out.Content + } + if out.IsError != nil { + ev.IsError = *out.IsError + acc.IsError = out.IsError + } + if out.Terminate != nil { + ev.Terminate = *out.Terminate + acc.Terminate = out.Terminate + } + }), +} + +type RunEndEvent struct { + SessionID string + TurnID string + Stop StopReason + Output string + Messages []Msg + MessageCounter int64 + Usage provider.Usage + Err error +} + +var RunEnd = Point[RunEndEvent, struct{}]{Kind: "run_end"} + +type SessionEvent struct { + SessionID string + ParentID string + AgentName string + Model string + Reason string +} + +var ( + SessionStart = Point[SessionEvent, struct{}]{Kind: "session_start"} + SessionEnd = Point[SessionEvent, struct{}]{Kind: "session_end"} +) + +type CompactEvent struct { + SessionID string + Trigger string + ContextTokens int + ContextWindow int +} + +type CancelResult struct { + Cancel bool + Reason string +} + +var BeforeCompact = Point[CompactEvent, CancelResult]{ + Kind: "before_compact", + Reduce: StopWhen[CompactEvent](func(r CancelResult) bool { return r.Cancel }), +} diff --git a/agent/hooks/reduce.go b/agent/hooks/reduce.go new file mode 100644 index 00000000..2a30110a --- /dev/null +++ b/agent/hooks/reduce.go @@ -0,0 +1,23 @@ +package hooks + +// StopWhen is the veto shape: the first handler whose result satisfies pred wins +// and the rest are skipped. Results that fail pred are discarded. +func StopWhen[E any, R any](pred func(R) bool) Reducer[E, R] { + return func(acc *R, _ *E, out R) bool { + if !pred(out) { + return false + } + *acc = out + return true + } +} + +// Fold is the mutation shape: apply merges each result into both the accumulator +// and the event, so the next handler sees what the previous one changed. It never +// short-circuits — every handler gets a turn. +func Fold[E any, R any](apply func(acc *R, ev *E, out R)) Reducer[E, R] { + return func(acc *R, ev *E, out R) bool { + apply(acc, ev, out) + return false + } +} diff --git a/agent/hooks_emit.go b/agent/hooks_emit.go new file mode 100644 index 00000000..6244beb2 --- /dev/null +++ b/agent/hooks_emit.go @@ -0,0 +1,175 @@ +package agent + +import ( + "context" + "fmt" + + "github.com/chainreactors/aiscan/agent/hooks" +) + +// The kernel reaches the typed hook registry only through these helpers. Each +// helper preserves the zero-handler fast path exposed by hooks.Registry. + +func runStartHook(ctx context.Context, cfg Config, systemPrompt string) (string, []ChatMessage) { + if !cfg.Hooks.Has(hooks.BeforeRun.Kind) { + return systemPrompt, nil + } + result, _ := hooks.BeforeRun.Emit(ctx, cfg.Hooks, hooks.RunStartEvent{ + SessionID: cfg.SessionID, + TurnID: cfg.TurnID, + AgentName: cfg.AgentName, + Model: cfg.Model, + SystemPrompt: systemPrompt, + ToolNames: toolNames(cfg), + }) + if result.SystemPrompt != nil { + systemPrompt = *result.SystemPrompt + } + return systemPrompt, result.Prepend +} + +func toolNames(cfg Config) []string { + if cfg.Tools == nil { + return nil + } + definitions := cfg.Tools.ToolDefinitions() + names := make([]string, 0, len(definitions)) + for _, definition := range definitions { + names = append(names, definition.Function.Name) + } + return names +} + +func transformContextHook(ctx context.Context, cfg Config, messages []ChatMessage, turn int) []ChatMessage { + if !cfg.Hooks.Has(hooks.Context.Kind) { + return messages + } + result, _ := hooks.Context.Emit(ctx, cfg.Hooks, hooks.ContextEvent{ + SessionID: cfg.SessionID, + Turn: turn, + Messages: messages, + }) + if result.Messages != nil { + return result.Messages + } + return messages +} + +// beforeTypedToolCall is fail-closed: a handler error means the call was not +// approved and is returned to the model as a tool error. +func beforeTypedToolCall(ctx context.Context, cfg Config, assistantMsg ChatMessage, tc ToolCall) toolExecution { + if !cfg.Hooks.Has(hooks.ToolCallHook.Kind) { + return toolExecution{} + } + decision, err := hooks.ToolCallHook.Emit(ctx, cfg.Hooks, hooks.ToolCallEvent{ + SessionID: cfg.SessionID, + TurnID: cfg.TurnID, + AssistantMessage: assistantMsg, + Call: tc, + SystemPrompt: cfg.SystemPrompt, + Messages: cfg.Messages, + }) + if err != nil { + return toolExecution{result: fmt.Sprintf("error: %s", err), isError: true, err: err} + } + if !decision.Block { + return toolExecution{} + } + reason := decision.Reason + if reason == "" { + reason = "tool execution was blocked" + } + return toolExecution{result: reason, isError: true} +} + +func afterTypedToolCall(ctx context.Context, cfg Config, tc ToolCall, execution toolExecution, durationMs int) toolExecution { + if !cfg.Hooks.Has(hooks.ToolResult.Kind) { + return execution + } + patch, err := hooks.ToolResult.Emit(ctx, cfg.Hooks, hooks.ToolResultEvent{ + SessionID: cfg.SessionID, + TurnID: cfg.TurnID, + Call: tc, + Content: execution.result, + IsError: execution.isError, + Terminate: execution.flow == ToolFlowTerminate, + DurationMs: durationMs, + Full: execution.fullResult, + }) + if err != nil { + execution.result = fmt.Sprintf("error: %s", err) + execution.isError = true + execution.err = err + return execution + } + if patch.Content != nil { + execution.result = *patch.Content + } + if patch.IsError != nil { + execution.isError = *patch.IsError + if !execution.isError { + execution.err = nil + } + } + if patch.Terminate != nil { + if *patch.Terminate { + execution.flow = ToolFlowTerminate + } else { + execution.flow = ToolFlowContinue + } + } + return execution +} + +func compactCanceled(ctx context.Context, cfg Config, trigger string, contextTokens int) (bool, string) { + if !cfg.Hooks.Has(hooks.BeforeCompact.Kind) { + return false, "" + } + result, _ := hooks.BeforeCompact.Emit(ctx, cfg.Hooks, hooks.CompactEvent{ + SessionID: cfg.SessionID, + Trigger: trigger, + ContextTokens: contextTokens, + ContextWindow: cfg.ContextWindow, + }) + return result.Cancel, result.Reason +} + +func emitRunEnd(ctx context.Context, cfg Config, result *Result) { + if result == nil || !cfg.Hooks.Has(hooks.RunEnd.Kind) { + return + } + _, _ = hooks.RunEnd.Emit(ctx, cfg.Hooks, hooks.RunEndEvent{ + SessionID: cfg.SessionID, + TurnID: cfg.TurnID, + Stop: result.Stop, + Output: result.Output, + Messages: result.Messages, + MessageCounter: result.MessageCounter, + Usage: result.TotalUsage, + Err: result.Err, + }) +} + +func emitSessionStart(ctx context.Context, cfg Config) { + if !cfg.Hooks.Has(hooks.SessionStart.Kind) { + return + } + _, _ = hooks.SessionStart.Emit(ctx, cfg.Hooks, sessionEvent(cfg, "")) +} + +func emitSessionEnd(ctx context.Context, cfg Config, reason string) { + if !cfg.Hooks.Has(hooks.SessionEnd.Kind) { + return + } + _, _ = hooks.SessionEnd.Emit(ctx, cfg.Hooks, sessionEvent(cfg, reason)) +} + +func sessionEvent(cfg Config, reason string) hooks.SessionEvent { + return hooks.SessionEvent{ + SessionID: cfg.SessionID, + ParentID: cfg.ParentSessionID, + AgentName: cfg.AgentName, + Model: cfg.Model, + Reason: reason, + } +} diff --git a/agent/inbox/context.go b/agent/inbox/context.go new file mode 100644 index 00000000..2e74db1a --- /dev/null +++ b/agent/inbox/context.go @@ -0,0 +1,23 @@ +package inbox + +import "context" + +type contextKey struct{} + +// ContextWithInbox scopes asynchronous command notifications to the agent +// session that invoked a tool. +func ContextWithInbox(ctx context.Context, ib Inbox) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, contextKey{}, ib) +} + +// FromContext returns the session inbox attached to ctx, if any. +func FromContext(ctx context.Context) Inbox { + if ctx == nil { + return nil + } + ib, _ := ctx.Value(contextKey{}).(Inbox) + return ib +} diff --git a/pkg/agent/inbox/expand.go b/agent/inbox/expand.go similarity index 97% rename from pkg/agent/inbox/expand.go rename to agent/inbox/expand.go index adaf45da..b2dedfb0 100644 --- a/pkg/agent/inbox/expand.go +++ b/agent/inbox/expand.go @@ -6,7 +6,7 @@ import ( "regexp" "strings" - "github.com/chainreactors/aiscan/pkg/agent/truncate" + "github.com/chainreactors/aiscan/core/truncate" ) const defaultMaxFileSize = truncate.DefaultMaxBytes diff --git a/pkg/agent/inbox/expand_test.go b/agent/inbox/expand_test.go similarity index 100% rename from pkg/agent/inbox/expand_test.go rename to agent/inbox/expand_test.go diff --git a/pkg/agent/inbox/inbox.go b/agent/inbox/inbox.go similarity index 100% rename from pkg/agent/inbox/inbox.go rename to agent/inbox/inbox.go diff --git a/pkg/agent/inbox/inbox_test.go b/agent/inbox/inbox_test.go similarity index 100% rename from pkg/agent/inbox/inbox_test.go rename to agent/inbox/inbox_test.go diff --git a/pkg/agent/inbox/message.go b/agent/inbox/message.go similarity index 98% rename from pkg/agent/inbox/message.go rename to agent/inbox/message.go index 6350eb25..c4297dba 100644 --- a/pkg/agent/inbox/message.go +++ b/agent/inbox/message.go @@ -5,7 +5,7 @@ import ( "strings" "time" - "github.com/chainreactors/aiscan/pkg/agent/provider" + "github.com/chainreactors/aiscan/agent/provider" ) type Origin string diff --git a/pkg/agent/input.go b/agent/input.go similarity index 98% rename from pkg/agent/input.go rename to agent/input.go index 4ed8c844..d07bfbc9 100644 --- a/pkg/agent/input.go +++ b/agent/input.go @@ -7,7 +7,7 @@ import ( "os" "strings" - "github.com/chainreactors/aiscan/pkg/aop" + "github.com/chainreactors/aiscan/core/aop" ) // maxInputImageBytes caps a single input image (20 MiB), matching common diff --git a/pkg/agent/input_test.go b/agent/input_test.go similarity index 98% rename from pkg/agent/input_test.go rename to agent/input_test.go index 623b410e..1f60fac5 100644 --- a/pkg/agent/input_test.go +++ b/agent/input_test.go @@ -7,7 +7,7 @@ import ( "strings" "testing" - "github.com/chainreactors/aiscan/pkg/aop" + "github.com/chainreactors/aiscan/core/aop" ) // pngBytes is a minimal PNG header so http.DetectContentType sniffs image/png. diff --git a/pkg/agent/loop.go b/agent/loop.go similarity index 88% rename from pkg/agent/loop.go rename to agent/loop.go index 744420e9..ee756a6c 100644 --- a/pkg/agent/loop.go +++ b/agent/loop.go @@ -9,14 +9,13 @@ import ( "sync" "time" + "github.com/chainreactors/aiscan/agent/inbox" + "github.com/chainreactors/aiscan/core/aop" + xcompact "github.com/chainreactors/aiscan/core/aop/x/compact" "github.com/chainreactors/aiscan/core/output" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/core/tool" - "github.com/chainreactors/aiscan/pkg/agent/inbox" - "github.com/chainreactors/aiscan/pkg/agent/truncate" - "github.com/chainreactors/aiscan/pkg/aop" - xcompact "github.com/chainreactors/aiscan/pkg/aop/x/compact" - "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/core/truncate" ) func runLoop(ctx context.Context, cfg Config) (*Result, error) { @@ -48,6 +47,7 @@ func runLoop(ctx context.Context, cfg Config) (*Result, error) { if result.Err != nil && stop == StopReasonError { em.errorEvt(result.Err, isRetryableError(result.Err)) } + emitRunEnd(ctx, cfg, result) if cfg.OnRunEnd != nil { cfg.OnRunEnd(result) } @@ -55,6 +55,10 @@ func runLoop(ctx context.Context, cfg Config) (*Result, error) { return result, err } + initialPrompt, prepend := runStartHook(ctx, cfg, cfg.SystemPrompt) + cfg.SystemPrompt = initialPrompt + transcript.append(prepend...) + for turn = 1; ; turn++ { if err := ctx.Err(); err != nil { failure := NewTextMessage("assistant", "") @@ -86,15 +90,15 @@ func runLoop(ctx context.Context, cfg Config) (*Result, error) { if cfg.SystemPromptFn != nil { systemPrompt = cfg.SystemPromptFn(&cfg) } - reqMessages := requestMessages(systemPrompt, transcript.messages, cfg.TransformContext) + reqMessages := requestMessages(ctx, cfg, systemPrompt, transcript.messages, turn) toolDefinitions := cfg.Tools.ToolDefinitions() contextTokens := transcript.estimatedContextTokens(estimateRequestTokens(reqMessages, toolDefinitions)) if shouldCompactContext(contextTokens, cfg.ContextWindow, cfg.Compaction) { - compacted, compactErr := runAutoCompaction(ctx, cfg, em, transcript, "threshold") + compacted, compactErr := runAutoCompaction(ctx, cfg, em, transcript, "threshold", contextTokens) if compactErr != nil { cfg.Logger.Warnf("auto-compaction failed: %s", compactErr) } else if compacted { - reqMessages = requestMessages(systemPrompt, transcript.messages, cfg.TransformContext) + reqMessages = requestMessages(ctx, cfg, systemPrompt, transcript.messages, turn) } } cfg.Logger.Debugf("[turn %d] sending %d messages to LLM", turn, len(reqMessages)) @@ -106,7 +110,7 @@ func runLoop(ctx context.Context, cfg Config) (*Result, error) { return end(nil, ctx.Err(), StopReasonCanceled) } if isContextOverflowError(err) && !overflowRecoveryAttempted { - compacted, compactErr := runAutoCompaction(ctx, cfg, em, transcript, "overflow") + compacted, compactErr := runAutoCompaction(ctx, cfg, em, transcript, "overflow", transcript.contextTokens) if compactErr != nil { cfg.Logger.Warnf("context overflow recovery failed: %s", compactErr) } else if compacted { @@ -121,7 +125,7 @@ func runLoop(ctx context.Context, cfg Config) (*Result, error) { assistantMsg = normalizeToolCalls(assistantMsg) if isLengthContextOverflow(assistantMsg.FinishReason, usage, cfg.ContextWindow) { if !overflowRecoveryAttempted { - compacted, compactErr := runAutoCompaction(ctx, cfg, em, transcript, "overflow") + compacted, compactErr := runAutoCompaction(ctx, cfg, em, transcript, "overflow", transcript.contextTokens) if compactErr != nil { cfg.Logger.Warnf("length overflow recovery failed: %s", compactErr) } else if compacted { @@ -266,11 +270,15 @@ func shouldCompactContext(contextTokens, contextWindow int, settings CompactionS return contextTokens > contextWindow-reserve } -func runAutoCompaction(ctx context.Context, cfg Config, em *aopEmitter, transcript *transcript, reason string) (bool, error) { +func runAutoCompaction(ctx context.Context, cfg Config, em *aopEmitter, transcript *transcript, reason string, contextTokens int) (bool, error) { reserve, keepRecent := effectiveCompactionLimits(cfg.ContextWindow, cfg.Compaction) if len(transcript.messages) < 2 || findCutPoint(transcript.messages, keepRecent) <= 0 { return false, nil } + if canceled, hookReason := compactCanceled(ctx, cfg, reason, contextTokens); canceled { + cfg.Logger.Debugf("compaction canceled by hook: %s", hookReason) + return false, nil + } em.status(xcompact.StateStart, "", nil) newMessages, result, err := compactHistory(ctx, CompactConfig{ @@ -501,9 +509,10 @@ type toolExecution struct { } func runToolCall(ctx context.Context, cfg Config, assistantMsg ChatMessage, tc ToolCall, turn int) toolExecution { + startedAt := time.Now() toolCtx := output.ContextWithCallID(ctx, tc.ID) toolCtx = withToolAgentConfig(toolCtx, cfg) - toolCtx = commands.ContextWithInbox(toolCtx, cfg.Inbox) + toolCtx = inbox.ContextWithInbox(toolCtx, cfg.Inbox) execution := beforeToolCall(toolCtx, cfg, assistantMsg, tc) if execution.result == "" && !execution.isError { toolResult, execErr := cfg.Tools.ExecuteTool(toolCtx, tc.Function.Name, tc.Function.Arguments) @@ -529,7 +538,7 @@ func runToolCall(ctx context.Context, cfg Config, assistantMsg ChatMessage, tc T "\n\n[truncated: showing %d/%d lines (%s of %s). Refine your query or use filter/parse tools to access specific parts.]", tr.OutputLines, tr.TotalLines, truncate.FormatSize(tr.OutputBytes), truncate.FormatSize(tr.TotalBytes)) } - return afterToolCall(toolCtx, cfg, assistantMsg, tc, execution) + return afterToolCall(toolCtx, cfg, assistantMsg, tc, execution, time.Since(startedAt).Milliseconds()) } // eventContent returns the AOP tool.result payload: a plain string, or the @@ -576,67 +585,65 @@ func toolResultToMessage(toolCallID string, exec toolExecution) ChatMessage { } func beforeToolCall(ctx context.Context, cfg Config, assistantMsg ChatMessage, tc ToolCall) toolExecution { - if cfg.BeforeToolCall == nil { - return toolExecution{} - } - before, err := cfg.BeforeToolCall(ctx, BeforeToolCallContext{ - AssistantMessage: assistantMsg, - ToolCall: tc, - SystemPrompt: cfg.SystemPrompt, - Messages: cfg.Messages, - }) - if err != nil { - return toolExecution{result: fmt.Sprintf("error: %s", err.Error()), isError: true, err: err} - } - if before == nil || !before.Block { - return toolExecution{} - } - result := before.Reason - if result == "" { - result = "tool execution was blocked" + if cfg.BeforeToolCall != nil { + before, err := cfg.BeforeToolCall(ctx, BeforeToolCallContext{ + AssistantMessage: assistantMsg, + ToolCall: tc, + SystemPrompt: cfg.SystemPrompt, + Messages: cfg.Messages, + }) + if err != nil { + return toolExecution{result: fmt.Sprintf("error: %s", err.Error()), isError: true, err: err} + } + if before != nil && before.Block { + result := before.Reason + if result == "" { + result = "tool execution was blocked" + } + return toolExecution{result: result, isError: true} + } } - return toolExecution{result: result, isError: true} + return beforeTypedToolCall(ctx, cfg, assistantMsg, tc) } -func afterToolCall(ctx context.Context, cfg Config, assistantMsg ChatMessage, tc ToolCall, execution toolExecution) toolExecution { - if cfg.AfterToolCall == nil { - return execution - } - after, err := cfg.AfterToolCall(ctx, AfterToolCallContext{ - AssistantMessage: assistantMsg, - ToolCall: tc, - Result: execution.result, - IsError: execution.isError, - SystemPrompt: cfg.SystemPrompt, - Messages: cfg.Messages, - }) - if err != nil { - execution.result = fmt.Sprintf("error: %s", err.Error()) - execution.isError = true - execution.err = err - return execution - } - if after == nil { - return execution - } - if after.Result != nil { - execution.result = *after.Result - } - if after.IsError != nil { - execution.isError = *after.IsError - if !execution.isError { - execution.err = nil +func afterToolCall(ctx context.Context, cfg Config, assistantMsg ChatMessage, tc ToolCall, execution toolExecution, durationMs int64) toolExecution { + if cfg.AfterToolCall != nil { + after, err := cfg.AfterToolCall(ctx, AfterToolCallContext{ + AssistantMessage: assistantMsg, + ToolCall: tc, + Result: execution.result, + IsError: execution.isError, + SystemPrompt: cfg.SystemPrompt, + Messages: cfg.Messages, + }) + if err != nil { + execution.result = fmt.Sprintf("error: %s", err.Error()) + execution.isError = true + execution.err = err + return execution + } + if after != nil { + if after.Result != nil { + execution.result = *after.Result + } + if after.IsError != nil { + execution.isError = *after.IsError + if !execution.isError { + execution.err = nil + } + } + execution.flow = after.Flow } } - execution.flow = after.Flow - return execution + return afterTypedToolCall(ctx, cfg, tc, execution, int(durationMs)) } -func requestMessages(systemPrompt string, messages []ChatMessage, transform TransformContextFunc) []ChatMessage { +func requestMessages(ctx context.Context, cfg Config, systemPrompt string, messages []ChatMessage, turn int) []ChatMessage { out := sanitizeMessages(append([]ChatMessage(nil), messages...)) - if transform != nil { - out = transform(out) + if cfg.TransformContext != nil { + out = cfg.TransformContext(out) } + out = transformContextHook(ctx, cfg, out, turn) if systemPrompt != "" { out = append([]ChatMessage{NewTextMessage("system", systemPrompt)}, out...) } diff --git a/agent/loop_context.go b/agent/loop_context.go new file mode 100644 index 00000000..5fdd5fbb --- /dev/null +++ b/agent/loop_context.go @@ -0,0 +1,29 @@ +package agent + +import "context" + +type loopSchedulerContextKey struct{} + +// ContextWithLoopScheduler scopes direct command execution to one runtime +// session. Agent tool calls carry the scheduler in their Config snapshot. +func ContextWithLoopScheduler(ctx context.Context, scheduler *LoopScheduler) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, loopSchedulerContextKey{}, scheduler) +} + +// LoopSchedulerFromContext resolves both direct-command and agent-tool-call +// contexts without exposing the agent's full Config. +func LoopSchedulerFromContext(ctx context.Context) *LoopScheduler { + if ctx == nil { + return nil + } + if scheduler, _ := ctx.Value(loopSchedulerContextKey{}).(*LoopScheduler); scheduler != nil { + return scheduler + } + if cfg, ok := toolAgentConfig(ctx); ok { + return cfg.LoopScheduler + } + return nil +} diff --git a/pkg/agent/loop_scheduler.go b/agent/loop_scheduler.go similarity index 98% rename from pkg/agent/loop_scheduler.go rename to agent/loop_scheduler.go index 03934976..cefead78 100644 --- a/pkg/agent/loop_scheduler.go +++ b/agent/loop_scheduler.go @@ -8,8 +8,8 @@ import ( "sync" "time" - "github.com/chainreactors/aiscan/pkg/agent/inbox" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/agent/inbox" + "github.com/chainreactors/aiscan/core/telemetry" ) type LoopMode int diff --git a/pkg/agent/loop_test.go b/agent/loop_test.go similarity index 98% rename from pkg/agent/loop_test.go rename to agent/loop_test.go index 8c6b2ac8..45decb4d 100644 --- a/pkg/agent/loop_test.go +++ b/agent/loop_test.go @@ -8,12 +8,12 @@ import ( "testing" "time" - "github.com/chainreactors/aiscan/pkg/agent/inbox" - "github.com/chainreactors/aiscan/pkg/agent/tmux" - "github.com/chainreactors/aiscan/pkg/agent/truncate" - "github.com/chainreactors/aiscan/pkg/aop" + "github.com/chainreactors/aiscan/agent/inbox" + "github.com/chainreactors/aiscan/agent/tmux" + "github.com/chainreactors/aiscan/core/aop" + "github.com/chainreactors/aiscan/core/telemetry" + "github.com/chainreactors/aiscan/core/truncate" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" ) func TestRunEmitsTurnEndAfterToolResults(t *testing.T) { @@ -1034,7 +1034,11 @@ func TestSessionCompletionInjectedIntoAgentLoop(t *testing.T) { t.Fatalf("Create: %v", err) } - time.Sleep(500 * time.Millisecond) + waitCtx, waitCancel := context.WithTimeout(context.Background(), 5*time.Second) + defer waitCancel() + if !ib.Wait(waitCtx) { + t.Fatal("timed out waiting for session completion") + } scripted := &scriptedProvider{ responses: []*ChatCompletionResponse{ @@ -1111,7 +1115,11 @@ func TestSessionCompletionMetadata(t *testing.T) { if err != nil { t.Fatalf("Create: %v", err) } - time.Sleep(500 * time.Millisecond) + waitCtx, waitCancel := context.WithTimeout(context.Background(), 5*time.Second) + defer waitCancel() + if !ib.Wait(waitCtx) { + t.Fatal("timed out waiting for session completion") + } received := ib.Drain() if len(received) == 0 { diff --git a/pkg/agent/overflow.go b/agent/overflow.go similarity index 100% rename from pkg/agent/overflow.go rename to agent/overflow.go diff --git a/pkg/agent/probe/llm.go b/agent/probe/llm.go similarity index 99% rename from pkg/agent/probe/llm.go rename to agent/probe/llm.go index c1d0ffa0..c87cb528 100644 --- a/pkg/agent/probe/llm.go +++ b/agent/probe/llm.go @@ -6,7 +6,7 @@ import ( "strings" "time" - "github.com/chainreactors/aiscan/pkg/agent" + "github.com/chainreactors/aiscan/agent" ) // LLMProbeRequest carries the connection parameters the user wants to verify diff --git a/pkg/agent/provider/anthropic.go b/agent/provider/anthropic.go similarity index 100% rename from pkg/agent/provider/anthropic.go rename to agent/provider/anthropic.go diff --git a/pkg/agent/provider/cache_test.go b/agent/provider/cache_test.go similarity index 100% rename from pkg/agent/provider/cache_test.go rename to agent/provider/cache_test.go diff --git a/pkg/agent/provider/capability_parity_test.go b/agent/provider/capability_parity_test.go similarity index 100% rename from pkg/agent/provider/capability_parity_test.go rename to agent/provider/capability_parity_test.go diff --git a/pkg/agent/provider/endpoint_hint_test.go b/agent/provider/endpoint_hint_test.go similarity index 100% rename from pkg/agent/provider/endpoint_hint_test.go rename to agent/provider/endpoint_hint_test.go diff --git a/pkg/agent/provider/errors.go b/agent/provider/errors.go similarity index 100% rename from pkg/agent/provider/errors.go rename to agent/provider/errors.go diff --git a/pkg/agent/provider/http.go b/agent/provider/http.go similarity index 100% rename from pkg/agent/provider/http.go rename to agent/provider/http.go diff --git a/pkg/agent/provider/openai.go b/agent/provider/openai.go similarity index 100% rename from pkg/agent/provider/openai.go rename to agent/provider/openai.go diff --git a/pkg/agent/provider/provider.go b/agent/provider/provider.go similarity index 100% rename from pkg/agent/provider/provider.go rename to agent/provider/provider.go diff --git a/pkg/agent/provider/provider_test.go b/agent/provider/provider_test.go similarity index 100% rename from pkg/agent/provider/provider_test.go rename to agent/provider/provider_test.go diff --git a/pkg/agent/provider/types.go b/agent/provider/types.go similarity index 100% rename from pkg/agent/provider/types.go rename to agent/provider/types.go diff --git a/pkg/agent/provider_swap_test.go b/agent/provider_swap_test.go similarity index 100% rename from pkg/agent/provider_swap_test.go rename to agent/provider_swap_test.go diff --git a/pkg/agent/retry.go b/agent/retry.go similarity index 98% rename from pkg/agent/retry.go rename to agent/retry.go index d1e2eaf8..bf61c0ca 100644 --- a/pkg/agent/retry.go +++ b/agent/retry.go @@ -12,9 +12,9 @@ import ( "strings" "time" - "github.com/chainreactors/aiscan/pkg/agent/provider" - "github.com/chainreactors/aiscan/pkg/aop" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/agent/provider" + "github.com/chainreactors/aiscan/core/aop" + "github.com/chainreactors/aiscan/core/telemetry" ) type imageDisabler interface { diff --git a/pkg/agent/retry_test.go b/agent/retry_test.go similarity index 89% rename from pkg/agent/retry_test.go rename to agent/retry_test.go index d8fa2f27..5ceea4c5 100644 --- a/pkg/agent/retry_test.go +++ b/agent/retry_test.go @@ -9,11 +9,11 @@ import ( "testing" "time" + "github.com/chainreactors/aiscan/agent/provider" + "github.com/chainreactors/aiscan/core/aop" "github.com/chainreactors/aiscan/core/eventbus" - "github.com/chainreactors/aiscan/pkg/agent/provider" - "github.com/chainreactors/aiscan/pkg/aop" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" ) func TestRetryOnTransientError(t *testing.T) { @@ -227,62 +227,6 @@ func TestStreamAssistantMessageReturnsContextErrorOnClosedCanceledStream(t *test } } -func TestProviderFailureDoesNotAutomaticallyFallback(t *testing.T) { - primary := &scriptedProvider{err: &APIError{StatusCode: 401, Message: "invalid api key"}} - fallback := &scriptedProvider{ - responses: []*ChatCompletionResponse{ - chatResponse(NewTextMessage("assistant", "from fallback")), - }, - } - - a := NewAgent(Config{ - Provider: primary, - Model: "primary-model", - Fallbacks: []ProviderEntry{{Provider: fallback, Model: "fallback-model"}}, - MaxRetries: 0, - Logger: telemetry.NopLogger(), - }) - - _, err := a.Run(context.Background(), TextInput("hello")) - if err == nil { - t.Fatal("Run() error = nil, want the primary provider error") - } - if len(fallback.requestsSnapshot()) != 0 { - t.Fatal("fallback provider was called automatically") - } -} - -func TestNoFallbackWhenPrimarySucceeds(t *testing.T) { - primary := &scriptedProvider{ - responses: []*ChatCompletionResponse{ - chatResponse(NewTextMessage("assistant", "from primary")), - }, - } - fallback := &scriptedProvider{ - responses: []*ChatCompletionResponse{ - chatResponse(NewTextMessage("assistant", "from fallback")), - }, - } - - a := NewAgent(Config{ - Provider: primary, - Fallbacks: []ProviderEntry{{Provider: fallback, Model: "fallback-model"}}, - MaxRetries: 0, - Logger: telemetry.NopLogger(), - }) - - result, err := a.Run(context.Background(), TextInput("hello")) - if err != nil { - t.Fatalf("Run() error = %v", err) - } - if result.Output != "from primary" { - t.Fatalf("Output = %q, want 'from primary'", result.Output) - } - if len(fallback.requestsSnapshot()) != 0 { - t.Fatal("fallback provider should not be called when primary succeeds") - } -} - // --- Image error recovery tests --- func TestImageErrorAutoRecovery(t *testing.T) { diff --git a/pkg/agent/session.go b/agent/session.go similarity index 100% rename from pkg/agent/session.go rename to agent/session.go diff --git a/pkg/agent/session_test.go b/agent/session_test.go similarity index 100% rename from pkg/agent/session_test.go rename to agent/session_test.go diff --git a/pkg/agent/subagent.go b/agent/subagent.go similarity index 98% rename from pkg/agent/subagent.go rename to agent/subagent.go index a6feb98d..b336a130 100644 --- a/pkg/agent/subagent.go +++ b/agent/subagent.go @@ -10,11 +10,11 @@ import ( "sync" "time" + "github.com/chainreactors/aiscan/agent/inbox" + "github.com/chainreactors/aiscan/core/aop/x/delegation" "github.com/chainreactors/aiscan/core/output" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/core/tool" - "github.com/chainreactors/aiscan/pkg/agent/inbox" - "github.com/chainreactors/aiscan/pkg/aop/x/delegation" - "github.com/chainreactors/aiscan/pkg/telemetry" ) type AgentType struct { diff --git a/pkg/agent/subagent_test.go b/agent/subagent_test.go similarity index 97% rename from pkg/agent/subagent_test.go rename to agent/subagent_test.go index c6752e85..6eb134ba 100644 --- a/pkg/agent/subagent_test.go +++ b/agent/subagent_test.go @@ -6,11 +6,11 @@ import ( "testing" "time" + "github.com/chainreactors/aiscan/agent/inbox" + "github.com/chainreactors/aiscan/core/aop" + "github.com/chainreactors/aiscan/core/aop/x/delegation" "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" - "github.com/chainreactors/aiscan/pkg/agent/inbox" - "github.com/chainreactors/aiscan/pkg/aop" - "github.com/chainreactors/aiscan/pkg/aop/x/delegation" "github.com/chainreactors/aiscan/pkg/commands" ) diff --git a/pkg/agent/tmux/manager.go b/agent/tmux/manager.go similarity index 100% rename from pkg/agent/tmux/manager.go rename to agent/tmux/manager.go diff --git a/pkg/agent/tmux/manager_test.go b/agent/tmux/manager_test.go similarity index 100% rename from pkg/agent/tmux/manager_test.go rename to agent/tmux/manager_test.go diff --git a/pkg/agent/tmux/process_alive_unix_test.go b/agent/tmux/process_alive_unix_test.go similarity index 100% rename from pkg/agent/tmux/process_alive_unix_test.go rename to agent/tmux/process_alive_unix_test.go diff --git a/pkg/agent/tmux/process_alive_windows_test.go b/agent/tmux/process_alive_windows_test.go similarity index 100% rename from pkg/agent/tmux/process_alive_windows_test.go rename to agent/tmux/process_alive_windows_test.go diff --git a/pkg/agent/tool_context.go b/agent/tool_context.go similarity index 100% rename from pkg/agent/tool_context.go rename to agent/tool_context.go diff --git a/pkg/agent/types.go b/agent/types.go similarity index 88% rename from pkg/agent/types.go rename to agent/types.go index 87ac0316..957c0164 100644 --- a/pkg/agent/types.go +++ b/agent/types.go @@ -5,13 +5,14 @@ import ( crand "crypto/rand" "encoding/hex" + "github.com/chainreactors/aiscan/agent/hooks" + "github.com/chainreactors/aiscan/agent/inbox" + "github.com/chainreactors/aiscan/agent/provider" + "github.com/chainreactors/aiscan/core/aop" + "github.com/chainreactors/aiscan/core/aop/x/delegation" "github.com/chainreactors/aiscan/core/eventbus" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/core/tool" - "github.com/chainreactors/aiscan/pkg/agent/inbox" - "github.com/chainreactors/aiscan/pkg/agent/provider" - "github.com/chainreactors/aiscan/pkg/aop" - "github.com/chainreactors/aiscan/pkg/aop/x/delegation" - "github.com/chainreactors/aiscan/pkg/telemetry" ) // Re-export provider types so external consumers only import agent. @@ -64,15 +65,17 @@ var ( // Agent-specific types. -type StopReason string +// StopReason is owned by agent/hooks so lifecycle events can carry it without +// introducing an import cycle back to the root agent package. +type StopReason = hooks.StopReason const ( - StopReasonCompleted StopReason = "completed" - StopReasonTerminated StopReason = "terminated" - StopReasonStopped StopReason = "stopped" - StopReasonBudget StopReason = "budget" - StopReasonError StopReason = "error" - StopReasonCanceled StopReason = "canceled" + StopReasonCompleted = hooks.StopReasonCompleted + StopReasonTerminated = hooks.StopReasonTerminated + StopReasonStopped = hooks.StopReasonStopped + StopReasonBudget = hooks.StopReasonBudget + StopReasonError = hooks.StopReasonError + StopReasonCanceled = hooks.StopReasonCanceled ) type TransformContextFunc func([]ChatMessage) []ChatMessage @@ -126,13 +129,9 @@ type CompactionSettings struct { } type Config struct { - Provider Provider - Tools tool.Executor - Model string - // Fallbacks is retained for source compatibility. It is deliberately - // ignored: provider selection is explicit and a run never switches models. - // Deprecated: configure and select provider profiles explicitly. - Fallbacks []ProviderEntry + Provider Provider + Tools tool.Executor + Model string SystemPrompt string SystemPromptFn SystemPromptFunc Messages []ChatMessage @@ -146,6 +145,9 @@ type Config struct { Logger telemetry.Logger TransformContext TransformContextFunc Bus *eventbus.Bus[aop.Event] + // Hooks is the typed extension registry shared by a runtime and its derived + // agents. Nil means no handlers and keeps the dispatch fast path allocation-free. + Hooks *hooks.Registry // OnRunEnd fires once per run with the final result — replaces the old // EventAgentEnd Messages subscription for session persistence. OnRunEnd func(*Result) @@ -197,6 +199,7 @@ func (c Config) WithCacheRetention(r CacheRetention) Config { c.CacheRetention = func (c Config) WithSessionID(id string) Config { c.SessionID = id; return c } func (c Config) WithTurnID(id string) Config { c.TurnID = id; return c } func (c Config) WithAgentName(name string) Config { c.AgentName = name; return c } +func (c Config) WithHooks(r *hooks.Registry) Config { c.Hooks = r; return c } func (c Config) WithOnRunEnd(fn func(*Result)) Config { c.OnRunEnd = fn; return c } func (c Config) WithLoopScheduler(s *LoopScheduler) Config { c.LoopScheduler = s diff --git a/cmd/agent/capability_test.go b/cmd/agent/capability_test.go new file mode 100644 index 00000000..97cf0968 --- /dev/null +++ b/cmd/agent/capability_test.go @@ -0,0 +1,24 @@ +package main + +import ( + "slices" + "testing" + + "github.com/chainreactors/aiscan/core/capability" + cfg "github.com/chainreactors/aiscan/core/config" +) + +func TestAgentCapabilitySetHasNoScanner(t *testing.T) { + want := []string{"arsenal", "core", "ioa"} + if got := capability.IDsSorted(); !slices.Equal(got, want) { + t.Fatalf("agent capabilities = %#v, want %#v", got, want) + } + for _, descriptor := range capability.All() { + if descriptor.Kind == capability.KindScanner { + t.Fatalf("agent linked scanner capability %q", descriptor.ID) + } + } + if got := cfg.CLICommandSummary(); got != "agent, web, serve" { + t.Fatalf("agent command summary = %q, want %q", got, "agent, web, serve") + } +} diff --git a/cmd/agent/main.go b/cmd/agent/main.go index ef5c5a6d..7a2517d6 100644 --- a/cmd/agent/main.go +++ b/cmd/agent/main.go @@ -10,14 +10,13 @@ import ( "time" cfg "github.com/chainreactors/aiscan/core/config" - transportpkg "github.com/chainreactors/aiscan/core/transport" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/core/telemetry" + "github.com/chainreactors/aiscan/pkg/runner" + transportpkg "github.com/chainreactors/aiscan/pkg/transport" goflags "github.com/jessevdk/go-flags" ) func main() { - cfg.ScannerEnabled = false - var option cfg.Option parser := goflags.NewParser(&option, goflags.Default&^goflags.PrintErrors) parser.Usage = `[OPTIONS] @@ -43,7 +42,7 @@ Examples: return } - cfgPath, err := cfg.ResolveRuntimeConfig(&option) + cfgPath, err := runner.ResolveRuntimeConfig(&option) if err != nil { fmt.Fprintf(os.Stderr, "error: %s\n", err) os.Exit(1) diff --git a/cmd/aiscan/capability_default_test.go b/cmd/aiscan/capability_default_test.go new file mode 100644 index 00000000..3209663e --- /dev/null +++ b/cmd/aiscan/capability_default_test.go @@ -0,0 +1,17 @@ +//go:build !full + +package main + +import ( + "slices" + "testing" + + "github.com/chainreactors/aiscan/core/capability" +) + +func TestDefaultCapabilitySet(t *testing.T) { + want := []string{"arsenal", "core", "gogo", "ioa", "neutron", "proton", "proxy", "scan", "search", "spray", "zombie"} + if got := capability.IDsSorted(); !slices.Equal(got, want) { + t.Fatalf("default capabilities = %#v, want %#v", got, want) + } +} diff --git a/cmd/aiscan/capability_full_test.go b/cmd/aiscan/capability_full_test.go new file mode 100644 index 00000000..025208a7 --- /dev/null +++ b/cmd/aiscan/capability_full_test.go @@ -0,0 +1,17 @@ +//go:build full + +package main + +import ( + "slices" + "testing" + + "github.com/chainreactors/aiscan/core/capability" +) + +func TestFullCapabilitySet(t *testing.T) { + want := []string{"arsenal", "browser", "core", "gogo", "ioa", "katana", "neutron", "passive", "proton", "proxy", "scan", "search", "spray", "zombie"} + if got := capability.IDsSorted(); !slices.Equal(got, want) { + t.Fatalf("full capabilities = %#v, want %#v", got, want) + } +} diff --git a/cmd/aiscan/cli.go b/cmd/aiscan/cli.go index 4a51d43d..82771790 100644 --- a/cmd/aiscan/cli.go +++ b/cmd/aiscan/cli.go @@ -15,9 +15,9 @@ import ( cfg "github.com/chainreactors/aiscan/core/config" "github.com/chainreactors/aiscan/core/output" - "github.com/chainreactors/aiscan/core/runner" - transportpkg "github.com/chainreactors/aiscan/core/transport" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/core/telemetry" + "github.com/chainreactors/aiscan/pkg/runner" + transportpkg "github.com/chainreactors/aiscan/pkg/transport" goflags "github.com/jessevdk/go-flags" ) @@ -137,7 +137,7 @@ func aiscan() { os.Exit(1) } - cfgPath, err := cfg.ResolveRuntimeConfig(&option) + cfgPath, err := runner.ResolveRuntimeConfig(&option) if err != nil { fmt.Fprintf(os.Stderr, "error: %s\n", err) os.Exit(1) diff --git a/cmd/aiscan/cli_test.go b/cmd/aiscan/cli_test.go index 631101ea..00c8b653 100644 --- a/cmd/aiscan/cli_test.go +++ b/cmd/aiscan/cli_test.go @@ -7,12 +7,12 @@ import ( "strings" "testing" + "github.com/chainreactors/aiscan/agent" cfg "github.com/chainreactors/aiscan/core/config" - "github.com/chainreactors/aiscan/core/runner" - "github.com/chainreactors/aiscan/pkg/agent" + "github.com/chainreactors/aiscan/pkg/runner" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" "github.com/chainreactors/aiscan/pkg/tui" "github.com/chainreactors/aiscan/skills" goflags "github.com/jessevdk/go-flags" @@ -145,9 +145,6 @@ func TestParseCLIRootTimeoutAppliesToAgent(t *testing.T) { } func TestDirectScannerModeSuppressesInitInfoByDefault(t *testing.T) { - if raceEnabled { - t.Skip("scanner pipeline has known races under -race; this test checks log output") - } var logBuf bytes.Buffer logger := telemetry.NewLogger(telemetry.LogConfig{Output: &logBuf}) err := runner.RunDirectScannerMode(context.Background(), &cfg.Option{ @@ -165,9 +162,6 @@ func TestDirectScannerModeSuppressesInitInfoByDefault(t *testing.T) { } func TestDirectScannerModeDebugShowsInitInfo(t *testing.T) { - if raceEnabled { - t.Skip("scanner pipeline has known races under -race; this test checks log output") - } var logBuf bytes.Buffer logger := telemetry.NewLogger(telemetry.LogConfig{Debug: true, Output: &logBuf}) err := runner.RunDirectScannerMode(context.Background(), &cfg.Option{ @@ -199,7 +193,7 @@ func TestParseCLIAgentAcceptsLLMFlags(t *testing.T) { if opt.BaseURL != "https://api.deepseek.com" || opt.APIKey != "KEY" || opt.Model != "deepseek-v4-pro" { t.Fatalf("llm options = %#v", opt.LLMOptions) } - pcfg := cfg.ProviderConfig(&opt) + pcfg := runner.ProviderConfig(&opt) if pcfg.Provider != "" { t.Fatalf("provider should be unresolved before agent.ResolveProvider, got %q", pcfg.Provider) } @@ -309,7 +303,7 @@ func TestParseCLIScanExtractsLLMFlags(t *testing.T) { if opt.AI || opt.APIKey != "KEY" || opt.Model != "deepseek-v4-pro" || opt.BaseURL != "https://api.deepseek.com" { t.Fatalf("llm options = %#v", opt.LLMOptions) } - pcfg := cfg.ProviderConfig(&opt) + pcfg := runner.ProviderConfig(&opt) if pcfg.Provider != "" { t.Fatalf("provider should be unresolved before agent.ResolveProvider, got %q", pcfg.Provider) } @@ -730,7 +724,7 @@ func TestAppConfigUsesCompiledDefaults(t *testing.T) { opt := &cfg.Option{} cfg.ApplyDefaults(opt) - appCfg := cfg.AppConfig(opt, cfg.RuntimeFeatures{ + appCfg := runner.AppConfig(opt, runner.RuntimeFeatures{ ProviderEnabled: true, ProviderOptional: true, AIEnabled: true, diff --git a/cmd/aiscan/race_norace_test.go b/cmd/aiscan/race_norace_test.go deleted file mode 100644 index 94aee521..00000000 --- a/cmd/aiscan/race_norace_test.go +++ /dev/null @@ -1,5 +0,0 @@ -//go:build !race - -package main - -const raceEnabled = false diff --git a/cmd/aiscan/race_test.go b/cmd/aiscan/race_test.go deleted file mode 100644 index ad0458fc..00000000 --- a/cmd/aiscan/race_test.go +++ /dev/null @@ -1,5 +0,0 @@ -//go:build race - -package main - -const raceEnabled = true diff --git a/cmd/aiscan/setup.go b/cmd/aiscan/setup.go index 87f16108..9485a8b0 100644 --- a/cmd/aiscan/setup.go +++ b/cmd/aiscan/setup.go @@ -7,16 +7,17 @@ import ( "os" "strings" + "github.com/chainreactors/aiscan/agent" + "github.com/chainreactors/aiscan/core/aop" + "github.com/chainreactors/aiscan/core/capability" cfg "github.com/chainreactors/aiscan/core/config" "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" "github.com/chainreactors/aiscan/core/pidlock" "github.com/chainreactors/aiscan/core/resources" - "github.com/chainreactors/aiscan/core/runner" - "github.com/chainreactors/aiscan/pkg/agent" - "github.com/chainreactors/aiscan/pkg/aop" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/pkg/runner" "github.com/chainreactors/aiscan/pkg/tui" "github.com/chainreactors/aiscan/skills" "github.com/chainreactors/aiscan/tools/scan" @@ -37,14 +38,14 @@ func init() { // Scanner engine initialization // --------------------------------------------------------------------------- -func scannerInit(ctx context.Context, a *runner.App, rc cfg.RuntimeConfig, logger telemetry.Logger) { +func scannerInit(ctx context.Context, a *runner.App, rc runner.ApplicationConfig, logger telemetry.Logger) { es := initEngines(ctx, rc.Scanner, logger) a.Engines = es registerScannerCommands(a.Commands, es, rc.Scanner, rc.Tools, a.Provider, a.ProviderConfig, a.Skills, a.DataBus, logger) } -func initEngines(ctx context.Context, sc cfg.ScannerConfig, logger telemetry.Logger) *engine.Set { +func initEngines(ctx context.Context, sc runner.ScannerConfig, logger telemetry.Logger) *engine.Set { engineSet, err := engine.InitWithOptions(ctx, resources.Options{ CyberhubURL: sc.CyberhubURL, APIKey: sc.CyberhubKey, @@ -67,8 +68,8 @@ func initEngines(ctx context.Context, sc cfg.ScannerConfig, logger telemetry.Log return engineSet } -func registerScannerCommands(cmdReg *commands.CommandRegistry, engineSet *engine.Set, scanCfg cfg.ScannerConfig, toolCfg cfg.ToolConfig, llmProvider agent.Provider, providerConfig agent.ProviderConfig, skillStore *skills.Store, dataBus *eventbus.Bus[output.ToolDataEvent], logger telemetry.Logger) { - var scanOpts []any +func registerScannerCommands(cmdReg *commands.CommandRegistry, engineSet *engine.Set, scanCfg runner.ScannerConfig, toolCfg runner.ToolConfig, llmProvider agent.Provider, providerConfig agent.ProviderConfig, skillStore *skills.Store, dataBus *eventbus.Bus[output.ToolDataEvent], logger telemetry.Logger) { + var scanOpts []scan.Option if scanCfg.AIEnabled && llmProvider != nil { scannerParent := agent.NewAgent(agent.Config{ Provider: llmProvider, @@ -99,19 +100,17 @@ func registerScannerCommands(cmdReg *commands.CommandRegistry, engineSet *engine WorkDir: workDir, BashTimeout: toolCfg.BashTimeout, SkillStore: skillStore, - EngineSet: engineSet, ScannerProxy: scanCfg.Proxy, - ScanOpts: scanOpts, Logger: logger, TavilyKeys: toolCfg.TavilyKeys, DataBus: dataBus, } + commands.Provide(deps, scan.OptsKey, scanOpts) if engineSet != nil { - deps.Resources = engineSet.Resources + commands.Provide(deps, engine.SetKey, engineSet) + commands.Provide(deps, resources.SetKey, engineSet.Resources) } - commands.BuildGroup("scanner", deps, cmdReg) - commands.BuildGroup("proxy", deps, cmdReg) - commands.BuildGroup("ioa", deps, cmdReg) + commands.BuildPlan(capability.Select(capability.Options{Groups: []string{"scanner", "proxy", "ioa"}}), deps, cmdReg) logger.Infof("%s", telemetry.StartupOK("scanner", strings.Join(cmdReg.GroupNames("scanner"), ","))) } diff --git a/cmd/aiscan/web_full.go b/cmd/aiscan/web_full.go index 4c75da50..bb300a0b 100644 --- a/cmd/aiscan/web_full.go +++ b/cmd/aiscan/web_full.go @@ -17,8 +17,8 @@ import ( "time" cfg "github.com/chainreactors/aiscan/core/config" - "github.com/chainreactors/aiscan/core/runner" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/core/telemetry" + "github.com/chainreactors/aiscan/pkg/runner" "github.com/chainreactors/aiscan/pkg/web" "github.com/chainreactors/aiscan/pkg/webproto" webstatic "github.com/chainreactors/aiscan/web" @@ -58,7 +58,7 @@ func runWeb(ctx context.Context, option, explicitOption *cfg.Option, opts webCom candidateOption = *explicitOption } candidateOption.ConfigFile = prepared.RuntimePath - if _, err := cfg.ResolveRuntimeConfigCandidate(&candidateOption); err != nil { + if _, err := runner.ResolveRuntimeConfigCandidate(&candidateOption); err != nil { return nil, err } candidate, err := initWebApp(ctx, &candidateOption, logger) @@ -194,7 +194,7 @@ func initWebApp(ctx context.Context, baseOption *cfg.Option, logger telemetry.Lo if baseOption != nil { option = *baseOption } - appCfg := cfg.AppConfig(&option, cfg.RuntimeFeatures{ + appCfg := runner.AppConfig(&option, runner.RuntimeFeatures{ ProviderEnabled: true, ProviderOptional: true, ToolsEnabled: true, diff --git a/cmd/aiscan/web_full_test.go b/cmd/aiscan/web_full_test.go index abe52acf..9e114ea0 100644 --- a/cmd/aiscan/web_full_test.go +++ b/cmd/aiscan/web_full_test.go @@ -11,7 +11,7 @@ import ( "testing" "github.com/chainreactors/aiscan/core/output" - "github.com/chainreactors/aiscan/core/runner" + "github.com/chainreactors/aiscan/pkg/runner" "github.com/chainreactors/aiscan/pkg/web" "github.com/chainreactors/aiscan/pkg/webproto" "gopkg.in/yaml.v3" diff --git a/cmd/runner/main.go b/cmd/runner/main.go index 0a5d4293..defca306 100644 --- a/cmd/runner/main.go +++ b/cmd/runner/main.go @@ -9,12 +9,14 @@ import ( "strings" "syscall" + "github.com/chainreactors/aiscan/core/capability" cfg "github.com/chainreactors/aiscan/core/config" "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" "github.com/chainreactors/aiscan/core/resources" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" + apprunner "github.com/chainreactors/aiscan/pkg/runner" "github.com/chainreactors/aiscan/pkg/webagent" "github.com/chainreactors/aiscan/tools/scan/engine" ) @@ -43,7 +45,7 @@ func main() { option := &cfg.Option{} option.ConfigFile = configFile - if _, err := cfg.ResolveRuntimeConfig(option); err != nil { + if _, err := apprunner.ResolveRuntimeConfig(option); err != nil { logger.Errorf("load config: %v", err) os.Exit(1) } @@ -90,17 +92,15 @@ func initTools(ctx context.Context, option *cfg.Option, logger telemetry.Logger, deps := &commands.Deps{ WorkDir: workDir, RunnerMode: true, - EngineSet: engineSet, Logger: logger, DataBus: dataBus, ScannerProxy: option.Proxy, } if engineSet != nil { - deps.Resources = engineSet.Resources - } - for _, group := range []string{"core", "scanner", "arsenal"} { - commands.BuildGroup(group, deps, registry) + commands.Provide(deps, engine.SetKey, engineSet) + commands.Provide(deps, resources.SetKey, engineSet.Resources) } + commands.BuildPlan(capability.Select(capability.Options{Groups: []string{"core", "scanner", "arsenal"}}), deps, registry) registry.SetLogger(logger) return registry, nil } diff --git a/cmd/runner/main_test.go b/cmd/runner/main_test.go index fb6a7609..d942021c 100644 --- a/cmd/runner/main_test.go +++ b/cmd/runner/main_test.go @@ -7,7 +7,7 @@ import ( cfg "github.com/chainreactors/aiscan/core/config" "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/core/telemetry" ) func TestInitToolsRegistersBash(t *testing.T) { diff --git a/pkg/aop/decode.go b/core/aop/decode.go similarity index 100% rename from pkg/aop/decode.go rename to core/aop/decode.go diff --git a/pkg/aop/event.go b/core/aop/event.go similarity index 100% rename from pkg/aop/event.go rename to core/aop/event.go diff --git a/pkg/aop/ext.go b/core/aop/ext.go similarity index 93% rename from pkg/aop/ext.go rename to core/aop/ext.go index 614abdb0..02cae0a2 100644 --- a/pkg/aop/ext.go +++ b/core/aop/ext.go @@ -9,7 +9,7 @@ import ( // // Ext/SetExt are the codec primitives for the extension map. Business code // must not call them directly — use the typed namespace packages under -// pkg/aop/x/ (or pkg/webproto for hub-owned namespaces) instead. +// core/aop/x/ (or pkg/webproto for hub-owned namespaces) instead. func Ext[T any](event Event, namespace string) (T, bool, error) { var value T raw, ok := event.Ext[namespace] diff --git a/pkg/aop/ext_types_gen.go b/core/aop/ext_types_gen.go similarity index 100% rename from pkg/aop/ext_types_gen.go rename to core/aop/ext_types_gen.go diff --git a/pkg/aop/gen_error.go b/core/aop/gen_error.go similarity index 100% rename from pkg/aop/gen_error.go rename to core/aop/gen_error.go diff --git a/pkg/aop/gen_message.go b/core/aop/gen_message.go similarity index 100% rename from pkg/aop/gen_message.go rename to core/aop/gen_message.go diff --git a/pkg/aop/gen_message_delta.go b/core/aop/gen_message_delta.go similarity index 100% rename from pkg/aop/gen_message_delta.go rename to core/aop/gen_message_delta.go diff --git a/pkg/aop/gen_session_start.go b/core/aop/gen_session_start.go similarity index 100% rename from pkg/aop/gen_session_start.go rename to core/aop/gen_session_start.go diff --git a/pkg/aop/gen_status.go b/core/aop/gen_status.go similarity index 100% rename from pkg/aop/gen_status.go rename to core/aop/gen_status.go diff --git a/pkg/aop/gen_tool_call.go b/core/aop/gen_tool_call.go similarity index 100% rename from pkg/aop/gen_tool_call.go rename to core/aop/gen_tool_call.go diff --git a/pkg/aop/gen_tool_result.go b/core/aop/gen_tool_result.go similarity index 100% rename from pkg/aop/gen_tool_result.go rename to core/aop/gen_tool_result.go diff --git a/pkg/aop/gen_turn.go b/core/aop/gen_turn.go similarity index 100% rename from pkg/aop/gen_turn.go rename to core/aop/gen_turn.go diff --git a/pkg/aop/gen_usage_session.go b/core/aop/gen_usage_session.go similarity index 100% rename from pkg/aop/gen_usage_session.go rename to core/aop/gen_usage_session.go diff --git a/pkg/aop/generate.go b/core/aop/generate.go similarity index 100% rename from pkg/aop/generate.go rename to core/aop/generate.go diff --git a/pkg/aop/schema_test.go b/core/aop/schema_test.go similarity index 100% rename from pkg/aop/schema_test.go rename to core/aop/schema_test.go diff --git a/pkg/aop/tool_result.go b/core/aop/tool_result.go similarity index 100% rename from pkg/aop/tool_result.go rename to core/aop/tool_result.go diff --git a/pkg/aop/tool_result_test.go b/core/aop/tool_result_test.go similarity index 100% rename from pkg/aop/tool_result_test.go rename to core/aop/tool_result_test.go diff --git a/pkg/aop/x/command/command.go b/core/aop/x/command/command.go similarity index 87% rename from pkg/aop/x/command/command.go rename to core/aop/x/command/command.go index 123965de..ba5e1365 100644 --- a/pkg/aop/x/command/command.go +++ b/core/aop/x/command/command.go @@ -1,6 +1,6 @@ package command -import "github.com/chainreactors/aiscan/pkg/aop" +import "github.com/chainreactors/aiscan/core/aop" const NS = "command" diff --git a/pkg/aop/x/compact/compact.go b/core/aop/x/compact/compact.go similarity index 86% rename from pkg/aop/x/compact/compact.go rename to core/aop/x/compact/compact.go index aa0987ab..1e77568a 100644 --- a/pkg/aop/x/compact/compact.go +++ b/core/aop/x/compact/compact.go @@ -1,6 +1,6 @@ package compact -import "github.com/chainreactors/aiscan/pkg/aop" +import "github.com/chainreactors/aiscan/core/aop" const ( NS = "compact" diff --git a/pkg/aop/x/compact/generate.go b/core/aop/x/compact/generate.go similarity index 100% rename from pkg/aop/x/compact/generate.go rename to core/aop/x/compact/generate.go diff --git a/pkg/aop/x/compact/types_gen.go b/core/aop/x/compact/types_gen.go similarity index 100% rename from pkg/aop/x/compact/types_gen.go rename to core/aop/x/compact/types_gen.go diff --git a/pkg/aop/x/delegation/delegation.go b/core/aop/x/delegation/delegation.go similarity index 83% rename from pkg/aop/x/delegation/delegation.go rename to core/aop/x/delegation/delegation.go index a4204828..0cd27e1b 100644 --- a/pkg/aop/x/delegation/delegation.go +++ b/core/aop/x/delegation/delegation.go @@ -1,6 +1,6 @@ package delegation -import "github.com/chainreactors/aiscan/pkg/aop" +import "github.com/chainreactors/aiscan/core/aop" const NS = "delegation" diff --git a/pkg/aop/x/delegation/generate.go b/core/aop/x/delegation/generate.go similarity index 100% rename from pkg/aop/x/delegation/generate.go rename to core/aop/x/delegation/generate.go diff --git a/pkg/aop/x/delegation/types_gen.go b/core/aop/x/delegation/types_gen.go similarity index 100% rename from pkg/aop/x/delegation/types_gen.go rename to core/aop/x/delegation/types_gen.go diff --git a/pkg/aop/x/eval/eval.go b/core/aop/x/eval/eval.go similarity index 90% rename from pkg/aop/x/eval/eval.go rename to core/aop/x/eval/eval.go index 237f9cbf..5b9c9188 100644 --- a/pkg/aop/x/eval/eval.go +++ b/core/aop/x/eval/eval.go @@ -1,6 +1,6 @@ package eval -import "github.com/chainreactors/aiscan/pkg/aop" +import "github.com/chainreactors/aiscan/core/aop" const ( NS = "eval" diff --git a/pkg/aop/x/eval/generate.go b/core/aop/x/eval/generate.go similarity index 100% rename from pkg/aop/x/eval/generate.go rename to core/aop/x/eval/generate.go diff --git a/pkg/aop/x/eval/types_gen.go b/core/aop/x/eval/types_gen.go similarity index 100% rename from pkg/aop/x/eval/types_gen.go rename to core/aop/x/eval/types_gen.go diff --git a/pkg/aop/x/ioa/generate.go b/core/aop/x/ioa/generate.go similarity index 100% rename from pkg/aop/x/ioa/generate.go rename to core/aop/x/ioa/generate.go diff --git a/pkg/aop/x/ioa/ioa.go b/core/aop/x/ioa/ioa.go similarity index 82% rename from pkg/aop/x/ioa/ioa.go rename to core/aop/x/ioa/ioa.go index a93a4a1b..2a231fa3 100644 --- a/pkg/aop/x/ioa/ioa.go +++ b/core/aop/x/ioa/ioa.go @@ -1,6 +1,6 @@ package ioa -import "github.com/chainreactors/aiscan/pkg/aop" +import "github.com/chainreactors/aiscan/core/aop" const NS = "ioa" diff --git a/pkg/aop/x/ioa/types_gen.go b/core/aop/x/ioa/types_gen.go similarity index 100% rename from pkg/aop/x/ioa/types_gen.go rename to core/aop/x/ioa/types_gen.go diff --git a/core/capability/capability.go b/core/capability/capability.go new file mode 100644 index 00000000..60c62d25 --- /dev/null +++ b/core/capability/capability.go @@ -0,0 +1,148 @@ +// Package capability is the single answer to "is this feature part of this +// binary". A capability exists if and only if its package is linked, and it is +// linked if and only if a blank import in cmd/*/imports*.go pulls it in — so +// build tags and blank imports stay the edition switch, and everything else +// (CLI help, scanner availability, skill gating, tool-group assembly) is +// derived from the descriptors registered here instead of from parallel global +// tables. +// +// The package deliberately imports nothing from aiscan so that core/config, +// skills, pkg/commands and pkg/tools can all depend on it. +package capability + +import ( + "sort" + "sync" +) + +type ID string + +type Kind uint8 + +const ( + KindTool Kind = iota // agent tool group + KindScanner // CLI-facing scanner command + KindService // ioa / proxy / web +) + +// Descriptor is what a capability package declares about itself in init(). +type Descriptor struct { + ID ID + Kind Kind + // Group is the command-factory group; empty means the ID is the group. + Group string + // CLIName is the top-level command name; empty means not CLI-facing. + CLIName string + // Summary is the word shown in the CLI command summary line. + Summary string + // UsageLine is the pre-aligned row shown in the scanner usage block. + UsageLine string + // Usage renders the command's full help lazily, so registering a + // capability never costs the work of building its usage text. + Usage func() string + // Skills lists skill names this capability unlocks. + Skills []string + // Optional marks a group selectable through --tools. + Optional bool + // Default enables an Optional group when --tools is empty. + Default bool + // Requires names the dependencies the factory needs, for the skip log. + Requires []string +} + +// Conflict records a duplicate registration. First registration wins; the +// duplicate is reported once at startup rather than silently shadowing. +type Conflict struct { + ID ID + Group string +} + +var ( + mu sync.RWMutex + order []ID + byID = map[ID]Descriptor{} + conflicts []Conflict +) + +// Register declares a capability. Called from init(); first registration wins. +func Register(d Descriptor) { + if d.ID == "" { + return + } + if d.Group == "" { + d.Group = string(d.ID) + } + mu.Lock() + defer mu.Unlock() + if _, exists := byID[d.ID]; exists { + conflicts = append(conflicts, Conflict{ID: d.ID, Group: d.Group}) + return + } + byID[d.ID] = d + order = append(order, d.ID) +} + +// All returns the descriptors in registration order. +func All() []Descriptor { + mu.RLock() + defer mu.RUnlock() + out := make([]Descriptor, 0, len(order)) + for _, id := range order { + out = append(out, byID[id]) + } + return out +} + +func Get(id ID) (Descriptor, bool) { + mu.RLock() + defer mu.RUnlock() + d, ok := byID[id] + return d, ok +} + +// Enabled reports whether the capability is linked into this binary. +func Enabled(id ID) bool { + mu.RLock() + defer mu.RUnlock() + _, ok := byID[id] + return ok +} + +func Conflicts() []Conflict { + mu.RLock() + defer mu.RUnlock() + return append([]Conflict(nil), conflicts...) +} + +// Groups lists every distinct factory group, in registration order. +func Groups() []string { + seen := map[string]bool{} + var out []string + for _, d := range All() { + if d.Group == "" || seen[d.Group] { + continue + } + seen[d.Group] = true + out = append(out, d.Group) + } + return out +} + +// IDsSorted is the stable identity of an edition, for golden tests. +func IDsSorted() []string { + ids := make([]string, 0, len(All())) + for _, d := range All() { + ids = append(ids, string(d.ID)) + } + sort.Strings(ids) + return ids +} + +// reset clears the registry. Tests only. +func reset() { + mu.Lock() + defer mu.Unlock() + order = nil + byID = map[ID]Descriptor{} + conflicts = nil +} diff --git a/core/capability/capability_test.go b/core/capability/capability_test.go new file mode 100644 index 00000000..b78c8c89 --- /dev/null +++ b/core/capability/capability_test.go @@ -0,0 +1,114 @@ +package capability + +import ( + "reflect" + "testing" +) + +func TestRegisterFirstWinsAndRecordsConflict(t *testing.T) { + reset() + Register(Descriptor{ID: "gogo", Kind: KindScanner, CLIName: "gogo", Summary: "gogo"}) + Register(Descriptor{ID: "gogo", Kind: KindScanner, CLIName: "shadow"}) + + d, ok := Get("gogo") + if !ok || d.CLIName != "gogo" { + t.Fatalf("first registration should win, got %#v (ok=%v)", d, ok) + } + if got := Conflicts(); len(got) != 1 || got[0].ID != "gogo" { + t.Fatalf("conflicts = %#v, want one entry for gogo", got) + } +} + +func TestGroupDefaultsToID(t *testing.T) { + reset() + Register(Descriptor{ID: "arsenal"}) + d, _ := Get("arsenal") + if d.Group != "arsenal" { + t.Fatalf("group = %q, want arsenal", d.Group) + } +} + +func TestQueriesOnlySeeLinkedCapabilities(t *testing.T) { + reset() + Register(Descriptor{ + ID: "gogo", Kind: KindScanner, Group: "scanner", + CLIName: "gogo", Summary: "gogo", UsageLine: " gogo Run gogo directly", + Usage: func() string { return "gogo help" }, + }) + + if !CLIAvailable("gogo") { + t.Fatal("gogo should be CLI-available") + } + if CLIAvailable("katana") { + t.Fatal("katana is not linked and must not be CLI-available") + } + if got := Summaries(); !reflect.DeepEqual(got, []string{"gogo"}) { + t.Fatalf("summaries = %#v", got) + } + if got := UsageLines(); !reflect.DeepEqual(got, []string{" gogo Run gogo directly"}) { + t.Fatalf("usage lines = %#v", got) + } + if usage, ok := Usage("gogo"); !ok || usage != "gogo help" { + t.Fatalf("usage = %q ok=%v", usage, ok) + } + if _, ok := Usage("katana"); ok { + t.Fatal("unlinked capability must not render usage") + } +} + +func TestSkillGatingFollowsLinkedCapability(t *testing.T) { + reset() + if SkillEnabled("katana") { + t.Fatal("katana skill must stay hidden while the capability is unlinked") + } + if !SkillEnabled("scan") { + t.Fatal("ungated skills are always enabled") + } + Register(Descriptor{ID: "katana", Kind: KindScanner, Skills: []string{"katana"}}) + if !SkillEnabled("katana") { + t.Fatal("katana skill should unlock once the capability is linked") + } +} + +func TestSelectHonoursOptionalAndDefault(t *testing.T) { + reset() + Register(Descriptor{ID: "core"}) + Register(Descriptor{ID: "search", Optional: true, Default: true}) + Register(Descriptor{ID: "browser", Optional: true, Default: true}) + Register(Descriptor{ID: "ioa", Optional: true}) + + plan := Select(Options{}) + for _, id := range []ID{"core", "search", "browser"} { + if !plan.Has(id) { + t.Fatalf("%s should be enabled by default", id) + } + } + if plan.Has("ioa") { + t.Fatal("non-default optional capability must stay off") + } + + plan = Select(Options{OptionalTools: []string{"browser"}}) + if plan.Has("search") { + t.Fatal("explicit --tools must not keep other optional capabilities") + } + if !plan.Has("browser") || !plan.Has("core") { + t.Fatal("explicit --tools must keep the selection and all non-optional capabilities") + } + + plan = Select(Options{Extra: []ID{"ioa"}}) + if !plan.Has("ioa") { + t.Fatal("Extra must force-enable a capability") + } +} + +func TestPlanGroupsFollowRegistrationOrder(t *testing.T) { + reset() + Register(Descriptor{ID: "core", Group: "core"}) + Register(Descriptor{ID: "gogo", Group: "scanner"}) + Register(Descriptor{ID: "spray", Group: "scanner"}) + Register(Descriptor{ID: "arsenal", Group: "arsenal"}) + + if got := Select(Options{}).Groups(); !reflect.DeepEqual(got, []string{"core", "scanner", "arsenal"}) { + t.Fatalf("groups = %#v", got) + } +} diff --git a/core/capability/gated.go b/core/capability/gated.go new file mode 100644 index 00000000..954488e9 --- /dev/null +++ b/core/capability/gated.go @@ -0,0 +1,29 @@ +package capability + +// gatedSkills maps a skill that ships in the embedded skill set to the +// capability that makes it usable. A skill is hidden unless its capability is +// linked — the embedded FS is the same in every edition, so this table (not a +// build tag) is what keeps the standard build from advertising skills whose +// tools it cannot run. +var gatedSkills = map[string]ID{ + "katana": "katana", + "passive": "passive", +} + +// SkillEnabled reports whether a skill should be visible in this binary. +func SkillEnabled(name string) bool { + id, gated := gatedSkills[name] + if !gated { + return true + } + return Enabled(id) +} + +// GatedSkills lists the skill names that depend on a capability being linked. +func GatedSkills() []string { + out := make([]string, 0, len(gatedSkills)) + for name := range gatedSkills { + out = append(out, name) + } + return out +} diff --git a/core/capability/plan.go b/core/capability/plan.go new file mode 100644 index 00000000..7d2c0799 --- /dev/null +++ b/core/capability/plan.go @@ -0,0 +1,78 @@ +package capability + +// A Plan is the set of capabilities one binary should actually assemble. It +// replaces the four divergent BuildGroup lists that each entry point used to +// keep, so "which groups does this binary build" has exactly one answer. + +type Options struct { + // Groups limits the plan to these assembly groups. Nil means every linked + // group. Entry points use this to describe their runtime surface without + // keeping private factory lists. + Groups []string + // OptionalTools is --tools / config tools. Empty selects the defaults. + OptionalTools []string + // Extra force-enables capabilities that become available at runtime, such + // as ioa once a client has connected. + Extra []ID +} + +type Plan struct { + enabled map[ID]bool + groups []string +} + +// Select resolves the registered descriptors against the caller's options. +// A capability that is not Optional is always part of the plan: it is linked, +// so it is meant to be there. +func Select(o Options) Plan { + groups := map[string]bool{} + for _, group := range o.Groups { + groups[group] = true + } + chosen := map[string]bool{} + for _, name := range o.OptionalTools { + chosen[name] = true + } + extra := map[ID]bool{} + for _, id := range o.Extra { + extra[id] = true + } + + p := Plan{enabled: map[ID]bool{}} + seen := map[string]bool{} + for _, d := range All() { + if len(groups) > 0 && !groups[d.Group] { + continue + } + switch { + case extra[d.ID]: + case !d.Optional: + case len(chosen) > 0: + if !chosen[string(d.ID)] && !chosen[d.Group] { + continue + } + case !d.Default: + continue + } + p.enabled[d.ID] = true + if d.Group != "" && !seen[d.Group] { + seen[d.Group] = true + p.groups = append(p.groups, d.Group) + } + } + return p +} + +func (p Plan) Has(id ID) bool { return p.enabled[id] } + +// Groups lists the factory groups to build, in registration order. +func (p Plan) Groups() []string { return append([]string(nil), p.groups...) } + +func (p Plan) HasGroup(group string) bool { + for _, g := range p.groups { + if g == group { + return true + } + } + return false +} diff --git a/core/capability/query.go b/core/capability/query.go new file mode 100644 index 00000000..e73ac837 --- /dev/null +++ b/core/capability/query.go @@ -0,0 +1,59 @@ +package capability + +// The queries here replace the Extra* globals that core/config used to carry: +// CLI availability, the usage block, the command summary and the lazy usage +// text all now come from the descriptors that are actually linked. + +// CLIAvailable reports whether name is a top-level CLI command in this binary. +func CLIAvailable(name string) bool { + _, ok := byCLIName(name) + return ok +} + +// UsageLines returns the pre-aligned usage rows of every linked CLI +// capability, in registration order. +func UsageLines() []string { + var out []string + for _, d := range All() { + if d.CLIName == "" || d.UsageLine == "" { + continue + } + out = append(out, d.UsageLine) + } + return out +} + +// Summaries returns the summary words of every linked CLI capability, in +// registration order. Callers prepend the binary's own modes (agent, web, …). +func Summaries() []string { + var out []string + for _, d := range All() { + if d.CLIName == "" || d.Summary == "" { + continue + } + out = append(out, d.Summary) + } + return out +} + +// Usage renders a CLI capability's help text, or false when the capability is +// not linked or declares no usage. +func Usage(name string) (string, bool) { + d, ok := byCLIName(name) + if !ok || d.Usage == nil { + return "", false + } + return d.Usage(), true +} + +func byCLIName(name string) (Descriptor, bool) { + if name == "" { + return Descriptor{}, false + } + for _, d := range All() { + if d.CLIName == name { + return d, true + } + } + return Descriptor{}, false +} diff --git a/core/config/defaults.go b/core/config/defaults.go new file mode 100644 index 00000000..804f1992 --- /dev/null +++ b/core/config/defaults.go @@ -0,0 +1,23 @@ +package config + +var ( + DefaultProvider = "openai" + DefaultBaseURL = "" + DefaultAPIKey = "" + DefaultModel = "" + + DefaultScannerProxy = "" + + DefaultCyberhubURL = "" + DefaultCyberhubKey = "" + DefaultCyberhubMode = "merge" + + DefaultVerify = "auto" + + DefaultIOAURL = "" + DefaultIOANodeID = "" + DefaultIOANodeName = "" + DefaultSpace = "" + + DefaultTavilyKeys = "" +) diff --git a/core/config/env.go b/core/config/env.go index 0dde5095..9517cab9 100644 --- a/core/config/env.go +++ b/core/config/env.go @@ -3,30 +3,20 @@ package config import ( "os" "strings" - - "github.com/chainreactors/aiscan/pkg/agent" ) type envLookup func(string) (string, bool) -func ResolveRuntimeConfig(option *Option) (string, error) { - return resolveRuntimeConfig(option, true) -} - -// ResolveRuntimeConfigCandidate resolves a staged configuration without -// mutating process-wide state. It is used to validate a Web reload candidate -// before the staged file is committed. -func ResolveRuntimeConfigCandidate(option *Option) (string, error) { - return resolveRuntimeConfig(option, false) -} - -func resolveRuntimeConfig(option *Option, applyProcessState bool) (string, error) { +// ResolveRuntimeConfig resolves parsed configuration with environment and +// defaults. Provider inference is supplied by the integration layer so config +// remains independent of concrete LLM implementations. +func ResolveRuntimeConfig(option *Option, applyProcessState bool, inferProvider func(string) string) (string, error) { explicit := *option configPath, err := LoadAndApplyConfig(option) if err != nil { return configPath, err } - applyEnvironment(option, explicit, os.LookupEnv) + applyEnvironment(option, explicit, os.LookupEnv, inferProvider) ApplyDefaults(option) if _, err := ResolveOutputPolicy(option); err != nil { return configPath, err @@ -37,19 +27,19 @@ func resolveRuntimeConfig(option *Option, applyProcessState bool) (string, error return configPath, nil } -func applyEnvironment(option *Option, explicit Option, lookup envLookup) { - applyLLMEnvironment(option, explicit, lookup) +func applyEnvironment(option *Option, explicit Option, lookup envLookup, inferProvider func(string) string) { + applyLLMEnvironment(option, explicit, lookup, inferProvider) applyScannerEnvironment(option, explicit, lookup) applyReconEnvironment(option, explicit, lookup) } -func applyLLMEnvironment(option *Option, explicit Option, lookup envLookup) { +func applyLLMEnvironment(option *Option, explicit Option, lookup envLookup, inferProvider func(string) string) { providerExplicit := strings.TrimSpace(explicit.Provider) != "" if v := firstEnv(lookup, "AISCAN_PROVIDER", "AISCAN_LLM_PROVIDER"); v != "" && !providerExplicit { option.Provider = v } - selectedProvider := selectedEnvProvider(option, lookup) + selectedProvider := selectedEnvProvider(option, lookup, inferProvider) if option.Provider == "" && selectedProvider != "" && !providerExplicit { option.Provider = selectedProvider } @@ -173,12 +163,12 @@ func applyReconEnvironment(option *Option, explicit Option, lookup envLookup) { } } -func selectedEnvProvider(option *Option, lookup envLookup) string { +func selectedEnvProvider(option *Option, lookup envLookup, inferProvider func(string) string) string { if v := strings.ToLower(strings.TrimSpace(option.Provider)); v != "" { return v } - if option.BaseURL != "" { - return agent.InferProviderFromBaseURL(option.BaseURL) + if option.BaseURL != "" && inferProvider != nil { + return inferProvider(option.BaseURL) } if firstEnv(lookup, "ANTHROPIC_API_KEY") != "" { return "anthropic" diff --git a/core/config/loader_test.go b/core/config/loader_test.go index 4c693bfe..b3b4af69 100644 --- a/core/config/loader_test.go +++ b/core/config/loader_test.go @@ -5,8 +5,6 @@ import ( "path/filepath" "strings" "testing" - - "github.com/chainreactors/aiscan/pkg/telemetry" ) func writeTestConfig(t *testing.T, dir, content string) string { @@ -306,9 +304,8 @@ search: if err := LoadConfig(filepath.Join(dir, "aiscan.yaml"), &option); err != nil { t.Fatal(err) } - cfg := AppConfig(&option, RuntimeFeatures{ToolsEnabled: true}, telemetry.NopLogger()) - if cfg.Tools.TavilyKeys != "K1,K2" { - t.Fatalf("tool config = %#v", cfg.Tools) + if option.SearchConfig.TavilyKeys != "K1,K2" { + t.Fatalf("search config = %#v", option.SearchConfig) } } @@ -323,7 +320,7 @@ scan: if err := LoadConfig(filepath.Join(dir, "aiscan.yaml"), &option); err != nil { t.Fatal(err) } - if got := AppConfig(&option, RuntimeFeatures{}, telemetry.NopLogger()).Scanner.VerifyMode; got != "critical" { + if got := option.ScanConfig.Verify; got != "critical" { t.Errorf("VerifyMode: got %q, want %q", got, "critical") } } @@ -517,7 +514,7 @@ cyberhub: defer os.Chdir(origDir) option := Option{} - if _, err := ResolveRuntimeConfig(&option); err != nil { + if _, err := ResolveRuntimeConfig(&option, true, testProviderInference); err != nil { t.Fatal(err) } @@ -556,7 +553,7 @@ llm: option.Model = "cli-model" option.BaseURL = "https://cli.example/v1" option.APIKey = "cli-key" - if _, err := ResolveRuntimeConfig(&option); err != nil { + if _, err := ResolveRuntimeConfig(&option, true, testProviderInference); err != nil { t.Fatal(err) } if option.Model != "cli-model" || option.BaseURL != "https://cli.example/v1" || option.APIKey != "cli-key" { @@ -581,7 +578,7 @@ llm: defer os.Chdir(origDir) option := Option{} - if _, err := ResolveRuntimeConfig(&option); err != nil { + if _, err := ResolveRuntimeConfig(&option, true, testProviderInference); err != nil { t.Fatal(err) } if option.Provider != "openai" || option.BaseURL != "https://openai-proxy.example/v1" || option.Model != "gpt-env" || option.APIKey != "openai-key" { @@ -606,7 +603,7 @@ llm: defer os.Chdir(origDir) option := Option{} - if _, err := ResolveRuntimeConfig(&option); err != nil { + if _, err := ResolveRuntimeConfig(&option, true, testProviderInference); err != nil { t.Fatal(err) } if option.Provider != "anthropic" || option.BaseURL != "https://anthropic-proxy.example/v1" || option.Model != "claude-env" || option.APIKey != "anthropic-key" { @@ -638,7 +635,7 @@ llm: defer os.Chdir(origDir) option := Option{} - if _, err := ResolveRuntimeConfig(&option); err != nil { + if _, err := ResolveRuntimeConfig(&option, true, testProviderInference); err != nil { t.Fatal(err) } if option.Model != "kimi-for-coding" { @@ -655,7 +652,7 @@ llm: defer os.Chdir(origDir) option := Option{} - if _, err := ResolveRuntimeConfig(&option); err != nil { + if _, err := ResolveRuntimeConfig(&option, true, testProviderInference); err != nil { t.Fatal(err) } if option.Model != "claude-opus-4-8" { @@ -687,7 +684,7 @@ llm: defer os.Chdir(origDir) option := Option{} - if _, err := ResolveRuntimeConfig(&option); err != nil { + if _, err := ResolveRuntimeConfig(&option, true, testProviderInference); err != nil { t.Fatal(err) } if option.BaseURL != "https://kiro.example/v1" { @@ -708,7 +705,7 @@ llm: defer os.Chdir(origDir) option := Option{} - if _, err := ResolveRuntimeConfig(&option); err != nil { + if _, err := ResolveRuntimeConfig(&option, true, testProviderInference); err != nil { t.Fatal(err) } if option.BaseURL != "https://borrowed.example/v1" { @@ -720,84 +717,6 @@ llm: }) } -func TestProvidersListOnly(t *testing.T) { - option := Option{} - option.Providers = []LLMProviderEntry{ - {Provider: "deepseek", APIKey: "key1", Model: "deepseek-chat"}, - {Provider: "openai", APIKey: "key2", Model: "gpt-4o"}, - } - - primary := ProviderConfig(&option) - if primary.Provider != "deepseek" || primary.APIKey != "key1" || primary.Model != "deepseek-chat" { - t.Errorf("primary should be providers[0], got %+v", primary) - } - - fallbacks := FallbackProviderConfigs(&option) - if len(fallbacks) != 1 || fallbacks[0].Provider != "openai" || fallbacks[0].Model != "gpt-4o" { - t.Errorf("fallback should be providers[1:], got %+v", fallbacks) - } -} - -func TestProvidersListWithSingleFields(t *testing.T) { - option := Option{} - option.Provider = "anthropic" - option.APIKey = "cli-key" - option.Providers = []LLMProviderEntry{ - {Provider: "deepseek", APIKey: "fb1", Model: "deepseek-chat"}, - } - - primary := ProviderConfig(&option) - if primary.Provider != "anthropic" || primary.APIKey != "cli-key" { - t.Errorf("single fields should win when set, got %+v", primary) - } - - fallbacks := FallbackProviderConfigs(&option) - if len(fallbacks) != 1 || fallbacks[0].Provider != "deepseek" { - t.Errorf("providers should be fallback when single fields set, got %+v", fallbacks) - } -} - -func TestProvidersListFromConfig(t *testing.T) { - dir := t.TempDir() - writeTestConfig(t, dir, ` -llm: - active_profile: openai - providers: - - id: deepseek - provider: deepseek - api_key: dk-111 - model: deepseek-chat - max_tokens: 8192 - context_window: 128000 - - id: openai - provider: openai - api_key: sk-222 - model: gpt-4o - max_tokens: 32768 - context_window: 1000000 -`) - - var opt Option - if err := LoadConfig(filepath.Join(dir, "aiscan.yaml"), &opt); err != nil { - t.Fatal(err) - } - if len(opt.Providers) != 2 { - t.Fatalf("expected 2 providers, got %d", len(opt.Providers)) - } - - primary := ProviderConfig(&opt) - if primary.Provider != "openai" || primary.APIKey != "sk-222" || - primary.MaxTokens != 32768 || primary.ContextWindow != 1000000 { - t.Errorf("primary from list: %+v", primary) - } - - fallbacks := FallbackProviderConfigs(&opt) - if len(fallbacks) != 1 || fallbacks[0].APIKey != "dk-111" || - fallbacks[0].MaxTokens != 8192 || fallbacks[0].ContextWindow != 128000 { - t.Errorf("fallbacks from list: %+v", fallbacks) - } -} - func TestResolveRuntimeConfigCandidateUsesStagedProfileAndExplicitCLIOverrides(t *testing.T) { for _, key := range []string{ "AISCAN_PROVIDER", "AISCAN_LLM_PROVIDER", "AISCAN_MODEL", "AISCAN_LLM_MODEL", @@ -824,25 +743,30 @@ llm: path := filepath.Join(dir, "aiscan.yaml") staged := Option{MiscOptions: MiscOptions{ConfigFile: path}} - if _, err := ResolveRuntimeConfigCandidate(&staged); err != nil { + if _, err := ResolveRuntimeConfig(&staged, false, testProviderInference); err != nil { t.Fatal(err) } - got := ProviderConfig(&staged) - if staged.ActiveProfile != "staged" || got.Provider != "openai" || got.Model != "staged-model" || got.APIKey != "staged-key" { - t.Fatalf("staged profile was not selected: option=%+v provider=%+v", staged.LLMOptions, got) + if staged.ActiveProfile != "staged" || len(staged.Providers) != 2 || staged.Providers[1].Model != "staged-model" { + t.Fatalf("staged profile was not loaded: option=%+v", staged.LLMOptions) } explicit := Option{ MiscOptions: MiscOptions{ConfigFile: path}, LLMOptions: LLMOptions{Provider: "deepseek", Model: "cli-model", APIKey: "cli-key"}, } - if _, err := ResolveRuntimeConfigCandidate(&explicit); err != nil { + if _, err := ResolveRuntimeConfig(&explicit, false, testProviderInference); err != nil { t.Fatal(err) } - got = ProviderConfig(&explicit) - if got.Provider != "deepseek" || got.Model != "cli-model" || got.APIKey != "cli-key" { - t.Fatalf("explicit CLI LLM values did not override staged config: %+v", got) + if explicit.Provider != "deepseek" || explicit.Model != "cli-model" || explicit.APIKey != "cli-key" { + t.Fatalf("explicit CLI LLM values did not override staged config: %+v", explicit.LLMOptions) + } +} + +func testProviderInference(baseURL string) string { + if strings.Contains(strings.ToLower(baseURL), "anthropic") { + return "anthropic" } + return "openai" } func withDefaults(t *testing.T, fn func()) { diff --git a/core/config/scanner.go b/core/config/scanner.go index 46553154..61ec269f 100644 --- a/core/config/scanner.go +++ b/core/config/scanner.go @@ -2,19 +2,9 @@ package config import ( "strings" -) - -var ExtraCommands = map[string]bool{} - -var ExtraUsageEntries []string - -var ExtraSummaryEntries []string -var ExtraScannerUsage = map[string]func() string{} - -// ScannerEnabled reports whether built-in scanner commands are available. -// Defaults to true; cmd/agent sets it to false. -var ScannerEnabled = true + "github.com/chainreactors/aiscan/core/capability" +) type ScannerCommands struct { Scan struct{} `command:"scan" description:"Run the scan pipeline"` @@ -28,47 +18,20 @@ type ScannerCommands struct { } func ScannerCommandAvailable(name string) bool { - if !ScannerEnabled { - return ExtraCommands[name] - } - switch name { - case "scan", "gogo", "spray", "zombie", "neutron": - return true - default: - return ExtraCommands[name] - } + return capability.CLIAvailable(name) } func ScannerUsageLines() string { - if !ScannerEnabled { - if len(ExtraUsageEntries) == 0 { - return "" - } - return strings.Join(ExtraUsageEntries, "\n") - } - base := ` gogo Run gogo directly - spray Run spray directly - zombie Run zombie directly - neutron Run neutron directly` - if len(ExtraUsageEntries) == 0 { - return base - } - return base + "\n" + strings.Join(ExtraUsageEntries, "\n") + return strings.Join(capability.UsageLines(), "\n") } func CLICommandSummary() string { - if !ScannerEnabled { - base := "agent, serve" - if len(ExtraSummaryEntries) == 0 { - return base - } - return base + ", " + strings.Join(ExtraSummaryEntries, ", ") - } - base := "agent, web, serve, scan, gogo, spray, zombie, neutron" - if len(ExtraSummaryEntries) == 0 { + base := "agent, web, serve" + summaries := capability.Summaries() + if len(summaries) == 0 { return base } - return base + ", " + strings.Join(ExtraSummaryEntries, ", ") + return base + ", " + strings.Join(summaries, ", ") } func IsScannerHelpRequest(args []string) bool { @@ -84,12 +47,5 @@ func IsScannerHelpRequest(args []string) bool { } func StaticScannerUsage(name string) (string, bool) { - if !ScannerCommandAvailable(name) { - return "", false - } - fn, ok := ExtraScannerUsage[name] - if !ok || fn == nil { - return "", false - } - return fn(), true + return capability.Usage(name) } diff --git a/core/config/scanner_katana.go b/core/config/scanner_katana.go deleted file mode 100644 index 8870151a..00000000 --- a/core/config/scanner_katana.go +++ /dev/null @@ -1,12 +0,0 @@ -//go:build full - -package config - -import katanacmd "github.com/chainreactors/aiscan/tools/katana" - -func init() { - ExtraCommands["katana"] = true - ExtraUsageEntries = append(ExtraUsageEntries, " katana Run katana web crawler") - ExtraSummaryEntries = append(ExtraSummaryEntries, "katana") - ExtraScannerUsage["katana"] = func() string { return katanacmd.New().Usage() } -} diff --git a/core/deps/architecture_test.go b/core/deps/architecture_test.go new file mode 100644 index 00000000..f582257c --- /dev/null +++ b/core/deps/architecture_test.go @@ -0,0 +1,154 @@ +package deps_test + +import ( + "go/parser" + "go/token" + "io/fs" + "os" + "path/filepath" + "strconv" + "strings" + "testing" +) + +const modulePath = "github.com/chainreactors/aiscan" + +func TestLayerImportsAreUnidirectional(t *testing.T) { + root := repositoryRoot(t) + assertNoFirstPartyImports(t, filepath.Join(root, "core"), map[string]bool{ + "agent": true, + "pkg": true, + "tools": true, + "cmd": true, + }) + assertNoFirstPartyImports(t, filepath.Join(root, "agent"), map[string]bool{ + "pkg": true, + "tools": true, + "cmd": true, + }) +} + +func TestLegacyPackagesCannotReturn(t *testing.T) { + root := repositoryRoot(t) + legacy := []struct { + dir string + importPath string + }{ + {dir: filepath.Join("pkg", "agent"), importPath: modulePath + "/pkg/" + "agent"}, + {dir: filepath.Join("pkg", "aop"), importPath: modulePath + "/pkg/" + "aop"}, + {dir: filepath.Join("pkg", "telemetry"), importPath: modulePath + "/pkg/" + "telemetry"}, + {dir: filepath.Join("pkg", "util"), importPath: modulePath + "/pkg/" + "util"}, + {dir: filepath.Join("core", "runner"), importPath: modulePath + "/core/" + "runner"}, + {dir: filepath.Join("core", "transport"), importPath: modulePath + "/core/" + "transport"}, + } + for _, item := range legacy { + legacyDir := filepath.Join(root, item.dir) + if _, err := os.Stat(legacyDir); err == nil { + t.Errorf("legacy package directory still exists: %s", legacyDir) + } else if !os.IsNotExist(err) { + t.Errorf("stat legacy package directory: %v", err) + } + } + + err := filepath.WalkDir(root, func(path string, entry fs.DirEntry, err error) error { + if err != nil { + return err + } + if entry.IsDir() { + if path != root && shouldSkipTree(root, path) { + return filepath.SkipDir + } + return nil + } + if filepath.Ext(path) != ".go" { + return nil + } + imports, parseErr := importsInFile(path) + if parseErr != nil { + return parseErr + } + for _, importPath := range imports { + for _, item := range legacy { + if importPath == item.importPath || strings.HasPrefix(importPath, item.importPath+"/") { + t.Errorf("legacy import %q in %s", importPath, relative(root, path)) + } + } + } + return nil + }) + if err != nil { + t.Fatal(err) + } +} + +func assertNoFirstPartyImports(t *testing.T, tree string, forbidden map[string]bool) { + t.Helper() + root := repositoryRoot(t) + err := filepath.WalkDir(tree, func(path string, entry fs.DirEntry, err error) error { + if err != nil { + return err + } + if entry.IsDir() || filepath.Ext(path) != ".go" || strings.HasSuffix(path, "_test.go") { + return nil + } + imports, parseErr := importsInFile(path) + if parseErr != nil { + return parseErr + } + for _, importPath := range imports { + if !strings.HasPrefix(importPath, modulePath+"/") { + continue + } + remainder := strings.TrimPrefix(importPath, modulePath+"/") + layer, _, _ := strings.Cut(remainder, "/") + if forbidden[layer] { + t.Errorf("forbidden dependency %q in %s", importPath, relative(root, path)) + } + } + return nil + }) + if err != nil { + t.Fatal(err) + } +} + +func importsInFile(path string) ([]string, error) { + file, err := parser.ParseFile(token.NewFileSet(), path, nil, parser.ImportsOnly) + if err != nil { + return nil, err + } + imports := make([]string, 0, len(file.Imports)) + for _, spec := range file.Imports { + value, err := strconv.Unquote(spec.Path.Value) + if err != nil { + return nil, err + } + imports = append(imports, value) + } + return imports, nil +} + +func repositoryRoot(t *testing.T) string { + t.Helper() + wd, err := os.Getwd() + if err != nil { + t.Fatal(err) + } + return filepath.Clean(filepath.Join(wd, "..", "..")) +} + +func shouldSkipTree(root, path string) bool { + rel := filepath.ToSlash(relative(root, path)) + return rel == ".git" || strings.HasPrefix(rel, ".git/") || + rel == "refer" || strings.HasPrefix(rel, "refer/") || + rel == "templates" || strings.HasPrefix(rel, "templates/") || + rel == "web/frontend/cyber-ui" || strings.HasPrefix(rel, "web/frontend/cyber-ui/") +} + +func relative(root, path string) string { + rel, err := filepath.Rel(root, path) + if err != nil { + return path + } + return filepath.ToSlash(rel) +} diff --git a/core/deps/deps.go b/core/deps/deps.go new file mode 100644 index 00000000..c9b19a20 --- /dev/null +++ b/core/deps/deps.go @@ -0,0 +1,72 @@ +// Package deps carries optional, typed dependencies from the assembly layer to +// the command factories. Keys are declared by the package that owns the value +// type, so the wiring layer (pkg/commands) never has to import — and therefore +// never links — the packages whose values it forwards. +package deps + +import "sync" + +// keyID gives a key its identity. Two keys declared with the same name are +// still distinct, and a key cannot be forged from a string. +type keyID struct{ name string } + +// Key names a dependency of type T. +type Key[T any] struct{ id *keyID } + +func NewKey[T any](name string) Key[T] { return Key[T]{id: &keyID{name: name}} } + +// Name reports the declared name, used when logging a missing dependency. +func Name[T any](k Key[T]) string { + if k.id == nil { + return "" + } + return k.id.name +} + +// Bag maps key identity to value. Accessors are package functions because Go +// methods cannot declare type parameters. +type Bag struct { + mu sync.RWMutex + values map[*keyID]any +} + +func New() *Bag { return &Bag{values: make(map[*keyID]any)} } + +// Set stores v under k. A nil Bag is a no-op so a half-built Deps cannot panic; +// use commands.Provide when the bag may not exist yet. +func Set[T any](b *Bag, k Key[T], v T) { + if b == nil || k.id == nil { + return + } + b.mu.Lock() + defer b.mu.Unlock() + if b.values == nil { + b.values = make(map[*keyID]any) + } + b.values[k.id] = v +} + +// Get returns the value stored under k. The key identity already guarantees the +// type; the assertion exists only because the backing map is untyped. +func Get[T any](b *Bag, k Key[T]) (T, bool) { + var zero T + if b == nil || k.id == nil { + return zero, false + } + b.mu.RLock() + raw, stored := b.values[k.id] + b.mu.RUnlock() + if !stored { + return zero, false + } + typed, ok := raw.(T) + if !ok { + return zero, false + } + return typed, true +} + +func Has[T any](b *Bag, k Key[T]) bool { + _, ok := Get(b, k) + return ok +} diff --git a/core/deps/quality_test.go b/core/deps/quality_test.go new file mode 100644 index 00000000..76abd8c5 --- /dev/null +++ b/core/deps/quality_test.go @@ -0,0 +1,286 @@ +package deps_test + +import ( + "bytes" + "encoding/json" + "fmt" + "go/ast" + "go/parser" + "go/token" + "io/fs" + "os" + "path/filepath" + "regexp" + "sort" + "strconv" + "strings" + "testing" +) + +type skipAllowance struct { + Path string `json:"path"` + Format string `json:"format"` + Count int `json:"count"` + Category string `json:"category"` + Reason string `json:"reason"` +} + +type skipKey struct { + Path string + Format string +} + +func TestSkipsMatchCentralRegistry(t *testing.T) { + root := repositoryRoot(t) + registryPath := filepath.Join(root, "test-skips.json") + data, err := os.ReadFile(registryPath) + if err != nil { + t.Fatal(err) + } + + var allowances []skipAllowance + if err := json.Unmarshal(data, &allowances); err != nil { + t.Fatalf("parse %s: %v", relative(root, registryPath), err) + } + + allowedCategories := map[string]bool{ + "capability": true, + "external_api": true, + "external_runtime": true, + "live_llm": true, + "platform": true, + } + want := make(map[skipKey]int, len(allowances)) + for _, allowance := range allowances { + key := skipKey{Path: filepath.ToSlash(allowance.Path), Format: allowance.Format} + switch { + case key.Path == "" || key.Format == "": + t.Errorf("skip registry entry must include path and format: %+v", allowance) + case allowance.Count <= 0: + t.Errorf("skip registry entry must have a positive count: %+v", allowance) + case !allowedCategories[allowance.Category]: + t.Errorf("skip registry entry has invalid category %q: %+v", allowance.Category, allowance) + case strings.TrimSpace(allowance.Reason) == "": + t.Errorf("skip registry entry must document its reason: %+v", allowance) + case want[key] != 0: + t.Errorf("duplicate skip registry entry for %s %q", key.Path, key.Format) + default: + want[key] = allowance.Count + } + } + + got, scanErrors := scanSkipCalls(root) + for _, scanErr := range scanErrors { + t.Error(scanErr) + } + + keys := make([]skipKey, 0, len(want)+len(got)) + seen := make(map[skipKey]bool, len(want)+len(got)) + for key := range want { + seen[key] = true + keys = append(keys, key) + } + for key := range got { + if !seen[key] { + keys = append(keys, key) + } + } + sort.Slice(keys, func(i, j int) bool { + if keys[i].Path == keys[j].Path { + return keys[i].Format < keys[j].Format + } + return keys[i].Path < keys[j].Path + }) + for _, key := range keys { + if got[key] != want[key] { + t.Errorf("skip registry mismatch for %s %q: found %d, registered %d", key.Path, key.Format, got[key], want[key]) + } + } +} + +func scanSkipCalls(root string) (map[skipKey]int, []error) { + got := make(map[skipKey]int) + var scanErrors []error + err := filepath.WalkDir(root, func(path string, entry fs.DirEntry, err error) error { + if err != nil { + return err + } + if entry.IsDir() { + if path != root && shouldSkipQualityTree(root, path) { + return filepath.SkipDir + } + return nil + } + ext := strings.ToLower(filepath.Ext(path)) + if ext == ".ts" || ext == ".tsx" || ext == ".js" || ext == ".jsx" { + calls, scriptErrors := scriptSkipCalls(root, path) + for key, count := range calls { + got[key] += count + } + scanErrors = append(scanErrors, scriptErrors...) + return nil + } + if ext != ".go" { + return nil + } + + file, parseErr := parser.ParseFile(token.NewFileSet(), path, nil, 0) + if parseErr != nil { + return parseErr + } + ast.Inspect(file, func(node ast.Node) bool { + call, ok := node.(*ast.CallExpr) + if !ok { + return true + } + selector, ok := call.Fun.(*ast.SelectorExpr) + if !ok || (selector.Sel.Name != "Skip" && selector.Sel.Name != "Skipf" && selector.Sel.Name != "SkipNow") { + return true + } + receiver, ok := selector.X.(*ast.Ident) + if !ok || (receiver.Name != "t" && receiver.Name != "b") { + return true + } + if len(call.Args) == 0 { + scanErrors = append(scanErrors, fmt.Errorf("unregistered reasonless skip in %s", relative(root, path))) + return true + } + literal, ok := call.Args[0].(*ast.BasicLit) + if !ok || literal.Kind != token.STRING { + scanErrors = append(scanErrors, fmt.Errorf("skip reason must be a string literal in %s", relative(root, path))) + return true + } + format, unquoteErr := strconv.Unquote(literal.Value) + if unquoteErr != nil { + scanErrors = append(scanErrors, fmt.Errorf("parse skip reason in %s: %w", relative(root, path), unquoteErr)) + return true + } + got[skipKey{Path: relative(root, path), Format: format}]++ + return true + }) + return nil + }) + if err != nil { + scanErrors = append(scanErrors, err) + } + return got, scanErrors +} + +var ( + scriptSkipStart = regexp.MustCompile(`\b(?:test|it|describe)\.skip\s*\(`) + scriptSkipReason = regexp.MustCompile("'[^']*'|\"[^\"]*\"|`[^`]*`") +) + +func scriptSkipCalls(root, path string) (map[skipKey]int, []error) { + data, err := os.ReadFile(path) + if err != nil { + return nil, []error{err} + } + rel := relative(root, path) + got := make(map[skipKey]int) + var scanErrors []error + for lineNumber, line := range strings.Split(string(data), "\n") { + starts := scriptSkipStart.FindAllStringIndex(line, -1) + for i, start := range starts { + end := len(line) + if i+1 < len(starts) { + end = starts[i+1][0] + } + literals := scriptSkipReason.FindAllString(line[start[0]:end], -1) + if len(literals) == 0 { + scanErrors = append(scanErrors, fmt.Errorf("skip reason must be a string literal in %s:%d", rel, lineNumber+1)) + continue + } + literal := literals[len(literals)-1] + reason := literal[1 : len(literal)-1] + got[skipKey{Path: rel, Format: reason}]++ + } + } + return got, scanErrors +} + +func TestRepositoryDebtMarkersCannotReturn(t *testing.T) { + root := repositoryRoot(t) + markers := [][]byte{[]byte("TO" + "DO"), []byte("FIX" + "ME")} + var failures []string + err := filepath.WalkDir(root, func(path string, entry fs.DirEntry, err error) error { + if err != nil { + return err + } + if entry.IsDir() { + if path != root && shouldSkipQualityTree(root, path) { + return filepath.SkipDir + } + return nil + } + + rel := relative(root, path) + if isBackupFile(entry.Name()) { + failures = append(failures, rel+": backup/editor artifact") + return nil + } + if !isDebtScannable(path) { + return nil + } + data, readErr := os.ReadFile(path) + if readErr != nil { + return readErr + } + for _, marker := range markers { + if bytes.Contains(data, marker) { + failures = append(failures, fmt.Sprintf("%s: contains forbidden debt marker %q", rel, marker)) + } + } + return nil + }) + if err != nil { + t.Fatal(err) + } + sort.Strings(failures) + for _, failure := range failures { + t.Error(failure) + } +} + +func shouldSkipQualityTree(root, path string) bool { + if shouldSkipTree(root, path) { + return true + } + rel := filepath.ToSlash(relative(root, path)) + if rel == "web/static" || strings.HasPrefix(rel, "web/static/") { + return true + } + parts := strings.Split(rel, "/") + for _, part := range parts { + switch part { + case "node_modules", "dist", "coverage", "playwright-report", "test-results", ".cache": + return true + } + } + return false +} + +func isBackupFile(name string) bool { + lower := strings.ToLower(name) + return strings.HasSuffix(lower, "~") || + strings.HasSuffix(lower, ".bak") || + strings.HasSuffix(lower, ".backup") || + strings.HasSuffix(lower, ".orig") || + strings.HasSuffix(lower, ".rej") || + strings.HasSuffix(lower, ".swp") || + strings.HasSuffix(lower, ".swo") || + strings.HasPrefix(lower, ".#") +} + +func isDebtScannable(path string) bool { + base := filepath.Base(path) + if base == "Makefile" || base == "Dockerfile" || base == ".gitattributes" || base == ".gitmodules" { + return true + } + switch strings.ToLower(filepath.Ext(path)) { + case ".css", ".go", ".html", ".js", ".json", ".jsx", ".md", ".mod", ".ps1", ".scss", ".sh", ".sum", ".toml", ".ts", ".tsx", ".yaml", ".yml": + return true + default: + return false + } +} diff --git a/core/harness/expect.go b/core/harness/expect.go deleted file mode 100644 index 4000695f..00000000 --- a/core/harness/expect.go +++ /dev/null @@ -1,204 +0,0 @@ -//go:build e2e - -package harness - -import ( - "encoding/json" - "fmt" - "strings" -) - -// ToolPattern describes an expected tool call. Built with the Tool() function -// and refined with chainable methods. -// -// Tool("bash").ArgContains("gogo").NoError() -// Tool("subagent").Action("create").Arg("name", "worker").Arg("mode", "async") -type ToolPattern struct { - tool string - action string - argChecks []argCheck - resultHas []string - resultNot []string - noError bool - isError bool - label string -} - -type argCheck struct { - key string - contains string -} - -func Tool(name string) ToolPattern { - return ToolPattern{tool: name, label: name} -} - -func (p ToolPattern) Action(action string) ToolPattern { - p.action = action - p.label = fmt.Sprintf("%s/%s", p.tool, action) - return p -} - -func (p ToolPattern) Arg(key, contains string) ToolPattern { - p.argChecks = append(p.argChecks, argCheck{key: key, contains: contains}) - return p -} - -func (p ToolPattern) ArgContains(substr string) ToolPattern { - p.argChecks = append(p.argChecks, argCheck{contains: substr}) - return p -} - -func (p ToolPattern) ResultHas(substr string) ToolPattern { - p.resultHas = append(p.resultHas, substr) - return p -} - -func (p ToolPattern) ResultNot(substr string) ToolPattern { - p.resultNot = append(p.resultNot, substr) - return p -} - -func (p ToolPattern) NoError() ToolPattern { - p.noError = true - return p -} - -func (p ToolPattern) IsError() ToolPattern { - p.isError = true - return p -} - -func (p ToolPattern) Label() string { return p.label } - -func (p ToolPattern) Match(e ToolExecution) bool { - if e.Name() != p.tool { - return false - } - if p.action != "" && !argsContainAction(e.Args(), p.action) { - return false - } - for _, ac := range p.argChecks { - if ac.key != "" { - if !argsFieldContains(e.Args(), ac.key, ac.contains) { - return false - } - } else { - if !strings.Contains(argsText(e.Args()), ac.contains) { - return false - } - } - } - for _, s := range p.resultHas { - if !strings.Contains(e.ResultText(), s) { - return false - } - } - for _, s := range p.resultNot { - if strings.Contains(e.ResultText(), s) { - return false - } - } - if p.noError && e.IsError() { - return false - } - if p.isError && !e.IsError() { - return false - } - return true -} - -func (p ToolPattern) describe() string { - var parts []string - parts = append(parts, p.tool) - if p.action != "" { - parts = append(parts, fmt.Sprintf("action=%s", p.action)) - } - for _, ac := range p.argChecks { - if ac.key != "" { - parts = append(parts, fmt.Sprintf("arg[%s]~%q", ac.key, ac.contains)) - } else { - parts = append(parts, fmt.Sprintf("args~%q", ac.contains)) - } - } - for _, s := range p.resultHas { - parts = append(parts, fmt.Sprintf("result~%q", s)) - } - return strings.Join(parts, " ") -} - -func argsContainAction(args any, action string) bool { - values, ok := args.(map[string]any) - if !ok { - return false - } - value, _ := values["action"].(string) - return value == action -} - -func argsFieldContains(args any, key, contains string) bool { - values, ok := args.(map[string]any) - if !ok { - return false - } - encoded, _ := json.Marshal(values[key]) - return strings.Contains(string(encoded), contains) -} - -// matchResult holds the result of matching expectations against actual tool calls. -type matchResult struct { - matched []matchPair - unmatched []ToolPattern -} - -type matchPair struct { - pattern ToolPattern - event ToolExecution - index int -} - -// matchUnordered finds a matching event for each pattern (greedy, unordered). -func matchUnordered(patterns []ToolPattern, events []ToolExecution) matchResult { - used := make([]bool, len(events)) - var matched []matchPair - var unmatched []ToolPattern - - for _, p := range patterns { - found := false - for i, e := range events { - if used[i] { - continue - } - if p.Match(e) { - matched = append(matched, matchPair{pattern: p, event: e, index: i}) - used[i] = true - found = true - break - } - } - if !found { - unmatched = append(unmatched, p) - } - } - return matchResult{matched: matched, unmatched: unmatched} -} - -// matchOrdered finds matching events in order (subsequence match). -func matchOrdered(patterns []ToolPattern, events []ToolExecution) matchResult { - var matched []matchPair - pi := 0 - for i, e := range events { - if pi >= len(patterns) { - break - } - if patterns[pi].Match(e) { - matched = append(matched, matchPair{pattern: patterns[pi], event: e, index: i}) - pi++ - } - } - var unmatched []ToolPattern - for _, p := range patterns[pi:] { - unmatched = append(unmatched, p) - } - return matchResult{matched: matched, unmatched: unmatched} -} diff --git a/core/harness/features_full_test.go b/core/harness/features_full_test.go deleted file mode 100644 index 932cd3cb..00000000 --- a/core/harness/features_full_test.go +++ /dev/null @@ -1,9 +0,0 @@ -//go:build e2e && full - -package harness - -func buildTags() string { return "emptytemplates noembed full" } - -func scannerHelpCommands() []string { - return []string{"gogo", "spray", "katana", "zombie", "neutron", "passive", "scan"} -} diff --git a/core/harness/features_test.go b/core/harness/features_test.go deleted file mode 100644 index 78cc3b81..00000000 --- a/core/harness/features_test.go +++ /dev/null @@ -1,9 +0,0 @@ -//go:build e2e && !full - -package harness - -func buildTags() string { return "emptytemplates noembed" } - -func scannerHelpCommands() []string { - return []string{"gogo", "spray", "zombie", "neutron", "scan"} -} diff --git a/core/harness/harness.go b/core/harness/harness.go deleted file mode 100644 index 03899bb4..00000000 --- a/core/harness/harness.go +++ /dev/null @@ -1,331 +0,0 @@ -//go:build e2e - -package harness - -import ( - "bytes" - "context" - "encoding/json" - "fmt" - "io" - "os" - "os/exec" - "path/filepath" - "runtime" - "strings" - "sync" - "testing" - "time" - - "github.com/chainreactors/aiscan/core/config" - "github.com/chainreactors/aiscan/pkg/aop" - "github.com/chainreactors/aiscan/pkg/webproto" -) - -var ( - cachedExe string - cachedExeOnce sync.Once - cachedExeErr error -) - -type Harness struct { - t *testing.T - exe string - workDir string - baseURL string - apiKey string - model string - timeout time.Duration - monitor *Monitor -} - -func (h *Harness) WithMonitor(out ...io.Writer) *Harness { - w := io.Writer(os.Stderr) - if len(out) > 0 { - w = out[0] - } - h.monitor = NewMonitor(w) - return h -} - -func New(t *testing.T) *Harness { - t.Helper() - - baseURL := os.Getenv("AISCAN_TEST_BASE_URL") - apiKey := os.Getenv("AISCAN_TEST_API_KEY") - model := os.Getenv("AISCAN_TEST_MODEL") - - if apiKey == "" { - t.Skip("AISCAN_TEST_API_KEY not set, skipping e2e test") - } - if baseURL == "" { - baseURL = "https://api.deepseek.com" - } - if model == "" { - model = "deepseek-v4-pro" - } - - cachedExeOnce.Do(func() { - cachedExe, cachedExeErr = buildOnce(t) - }) - if cachedExeErr != nil { - t.Fatalf("build aiscan: %v", cachedExeErr) - } - - h := &Harness{ - t: t, - exe: cachedExe, - workDir: t.TempDir(), - baseURL: baseURL, - apiKey: apiKey, - model: model, - timeout: 180 * time.Second, - } - if os.Getenv("AISCAN_MONITOR") != "" { - h.monitor = NewMonitor(os.Stderr) - } - return h -} - -func buildOnce(t *testing.T) (string, error) { - t.Helper() - dir, err := os.MkdirTemp("", "aiscan-e2e-*") - if err != nil { - return "", err - } - exeName := "aiscan-e2e" - if runtime.GOOS == "windows" { - exeName += ".exe" - } - exe := filepath.Join(dir, exeName) - args := []string{"build", "-tags", buildTags(), "-o", exe, "./cmd/aiscan"} - cmd := exec.Command("go", args...) - cmd.Dir = repoRoot(t) - out, err := cmd.CombinedOutput() - if err != nil { - return "", fmt.Errorf("%v\n%s", err, out) - } - return exe, nil -} - -func (h *Harness) llmArgs() []string { - return []string{ - "--base-url", h.baseURL, - "--api-key", h.apiKey, - "--model", h.model, - } -} - -func (h *Harness) Run(args ...string) *RunResult { - h.t.Helper() - return h.RunWithTimeout(h.timeout, args...) -} - -func (h *Harness) RunWithTimeout(timeout time.Duration, args ...string) *RunResult { - h.t.Helper() - - var fullArgs []string - switch { - case len(args) > 0 && args[0] == "agent": - fullArgs = h.agentCLIArgs(args[1:]...) - case len(args) == 1 && args[0] == "--version": - fullArgs = []string{"--no-color", "--quiet", "--version"} - default: - fullArgs = append(h.llmArgs(), "--no-color", "--quiet") - fullArgs = append(fullArgs, args...) - } - - ctx, cancel := context.WithTimeout(context.Background(), timeout) - defer cancel() - - cmd := exec.CommandContext(ctx, h.exe, fullArgs...) - cmd.Dir = h.workDir - - var stdout, stderr bytes.Buffer - cmd.Stdout = &stdout - cmd.Stderr = &stderr - - start := time.Now() - err := cmd.Run() - duration := time.Since(start) - - exitCode := processExitCode(err) - - result := &RunResult{ - Stdout: stdout.String(), - Stderr: stderr.String(), - ExitCode: exitCode, - Duration: duration, - } - - h.t.Logf("ran: aiscan %s (exit=%d, duration=%s, turns=%d, tools=%d)", - strings.Join(args, " "), exitCode, duration.Round(time.Millisecond), - result.Turns(), len(result.ToolCalls())) - if exitCode != 0 { - h.t.Logf("stderr: %s", clip(stderr.String(), 2000)) - } - - return result -} - -func (h *Harness) WorkFile(name string) string { - return filepath.Join(h.workDir, name) -} - -// --- convenience runners --- - -func (h *Harness) Agent(prompt string, extraArgs ...string) *RunResult { - h.t.Helper() - return h.AgentWithTimeout(h.timeout, prompt, extraArgs...) -} - -func (h *Harness) AgentWithTimeout(timeout time.Duration, prompt string, extraArgs ...string) *RunResult { - h.t.Helper() - - fullArgs := h.agentCLIArgs(extraArgs...) - fullArgs = append(fullArgs, "--transport", "stdio") - - ctx, cancel := context.WithTimeout(context.Background(), timeout) - defer cancel() - - cmd := exec.CommandContext(ctx, h.exe, fullArgs...) - cmd.Dir = h.workDir - stdin, err := cmd.StdinPipe() - if err != nil { - h.t.Fatalf("agent stdin pipe: %v", err) - } - stdout, err := cmd.StdoutPipe() - if err != nil { - h.t.Fatalf("agent stdout pipe: %v", err) - } - var stderr bytes.Buffer - cmd.Stderr = &stderr - - start := time.Now() - if err := cmd.Start(); err != nil { - h.t.Fatalf("start aiscan agent: %v", err) - } - - encoder := json.NewEncoder(stdin) - openErr := encoder.Encode(webproto.Message{ - Type: webproto.TypeSessionOpen, - Payload: webproto.MustJSON(webproto.SessionOpenPayload{SessionID: "harness"}), - }) - writeErr := openErr - if writeErr == nil { - writeErr = encoder.Encode(webproto.Message{ - Type: webproto.TypeRun, TurnID: "turn-1", - Payload: webproto.MustJSON(webproto.RunPayload{ - SessionID: "harness", - Parts: []aop.MessagePart{{Type: aop.PartText, Text: prompt}}, - }), - }) - } - closeErr := stdin.Close() - if writeErr != nil || closeErr != nil { - cancel() - } - - output, events, streamErr := consumeAgentStream(stdout, h.monitor) - if streamErr != nil { - cancel() - } - waitErr := cmd.Wait() - duration := time.Since(start) - - exitCode := processExitCode(waitErr) - if writeErr != nil { - exitCode = -1 - fmt.Fprintf(&stderr, "write stdio request: %v\n", writeErr) - } else if closeErr != nil { - exitCode = -1 - fmt.Fprintf(&stderr, "close stdio request: %v\n", closeErr) - } - if streamErr != nil { - exitCode = -1 - fmt.Fprintf(&stderr, "read AOP stdout: %v\n", streamErr) - } - - result := &RunResult{ - Stdout: output, - Stderr: stderr.String(), - ExitCode: exitCode, - Duration: duration, - Events: events, - } - - h.t.Logf("ran: aiscan agent (exit=%d, duration=%s, turns=%d, tools=%d)", - exitCode, duration.Round(time.Millisecond), result.Turns(), len(result.ToolCalls())) - if exitCode != 0 { - h.t.Logf("stderr: %s", clip(stderr.String(), 2000)) - } - - return result -} - -func (h *Harness) AgentWithInput(prompt string, inputs []string, extraArgs ...string) *RunResult { - h.t.Helper() - args := make([]string, 0, len(inputs)*2+len(extraArgs)) - for _, input := range inputs { - args = append(args, "-i", input) - } - args = append(args, extraArgs...) - task := fmt.Sprintf("%s\n\nTargets:\n%s", prompt, config.FormatInputs(inputs)) - return h.Agent(task, args...) -} - -func (h *Harness) agentCLIArgs(extraArgs ...string) []string { - args := []string{"--no-color", "--quiet", "agent"} - args = append(args, h.llmArgs()...) - return append(args, extraArgs...) -} - -func (h *Harness) Scanner(name string, scannerArgs ...string) *RunResult { - h.t.Helper() - args := []string{name} - args = append(args, scannerArgs...) - return h.Run(args...) -} - -func (h *Harness) ScannerAI(name string, scannerArgs ...string) *RunResult { - h.t.Helper() - args := []string{"--ai", name} - args = append(args, scannerArgs...) - return h.Run(args...) -} - -// --- helpers --- - -func repoRoot(t *testing.T) string { - t.Helper() - wd, err := os.Getwd() - if err != nil { - t.Fatal(err) - } - return filepath.Clean(filepath.Join(wd, "..", "..")) -} - -func envOrDefault(key, fallback string) string { - if v := os.Getenv(key); v != "" { - return v - } - return fallback -} - -func processExitCode(err error) int { - if err == nil { - return 0 - } - if exitErr, ok := err.(*exec.ExitError); ok { - return exitErr.ExitCode() - } - return -1 -} - -func clip(s string, maxLen int) string { - s = strings.TrimSpace(s) - if len(s) <= maxLen { - return s - } - return s[:maxLen] + "... (truncated)" -} diff --git a/core/harness/harness_test.go b/core/harness/harness_test.go deleted file mode 100644 index 04e1afdc..00000000 --- a/core/harness/harness_test.go +++ /dev/null @@ -1,1528 +0,0 @@ -//go:build e2e - -package harness - -import ( - "bytes" - "context" - "encoding/json" - "fmt" - "net/http/httptest" - "os" - "os/exec" - "strings" - "testing" - "time" - - "github.com/chainreactors/aiscan/pkg/aop" - "github.com/chainreactors/aiscan/pkg/webproto" - ioaclient "github.com/chainreactors/ioa/client" - "github.com/chainreactors/ioa/protocols" - ioaserver "github.com/chainreactors/ioa/server" -) - -// ===================================================================== -// init -// ===================================================================== - -func init() { - if _, err := exec.LookPath("go"); err != nil { - panic("go compiler not found; e2e tests require Go toolchain") - } -} - -// ===================================================================== -// Agent — basic prompts and tool use -// ===================================================================== - -func TestAgentSimplePrompt(t *testing.T) { - h := New(t) - Intent{ - Name: "simple-prompt", - Prompt: "What is 2+2? Reply with just the number.", - OutputContains: []string{"4"}, - MaxTurns: 2, - JudgeCriteria: "The agent must reply with the number 4. No tool calls needed. The answer must be mathematically correct.", - }.Run(t, h) -} - -// TestAgentDualSessionInterleaved drives the stdio host with two concurrent -// sessions: messages for sess-a and sess-b are written interleaved and both -// sessions must run to completion independently. -func TestAgentDualSessionInterleaved(t *testing.T) { - h := New(t) - - fullArgs := h.agentCLIArgs() - fullArgs = append(fullArgs, "--transport", "stdio") - - ctx, cancel := context.WithTimeout(context.Background(), h.timeout) - defer cancel() - - cmd := exec.CommandContext(ctx, h.exe, fullArgs...) - cmd.Dir = h.workDir - stdin, err := cmd.StdinPipe() - if err != nil { - t.Fatalf("stdin pipe: %v", err) - } - stdout, err := cmd.StdoutPipe() - if err != nil { - t.Fatalf("stdout pipe: %v", err) - } - var stderr bytes.Buffer - cmd.Stderr = &stderr - if err := cmd.Start(); err != nil { - t.Fatalf("start agent: %v", err) - } - - encoder := json.NewEncoder(stdin) - writeFrame := func(message webproto.Message) { - t.Helper() - if err := encoder.Encode(message); err != nil { - t.Fatalf("write %s: %v", message.Type, err) - } - } - writeRun := func(sessionID, turnID, text string) { - writeFrame(webproto.Message{ - Type: webproto.TypeRun, TurnID: turnID, - Payload: webproto.MustJSON(webproto.RunPayload{ - SessionID: sessionID, - Parts: []aop.MessagePart{{Type: aop.PartText, Text: text}}, - }), - }) - } - writeFrame(webproto.Message{Type: webproto.TypeSessionOpen, Payload: webproto.MustJSON(webproto.SessionOpenPayload{SessionID: "sess-a"})}) - writeFrame(webproto.Message{Type: webproto.TypeSessionOpen, Payload: webproto.MustJSON(webproto.SessionOpenPayload{SessionID: "sess-b"})}) - writeRun("sess-a", "turn-a1", "Reply with exactly: ALPHA") - writeRun("sess-b", "turn-b1", "Reply with exactly: BRAVO") - writeRun("sess-a", "turn-a2", "Reply with exactly: ALPHA2") - if err := stdin.Close(); err != nil { - t.Fatalf("close stdin: %v", err) - } - - type sessionState struct { - ended int - wantEnd int - outputs []string - } - sessions := map[string]*sessionState{"sess-a": {wantEnd: 2}, "sess-b": {wantEnd: 1}} - decoder := json.NewDecoder(stdout) - for { - var message webproto.Message - if err := decoder.Decode(&message); err != nil { - break - } - if message.Type != webproto.TypeAOP { - continue - } - var event aop.Event - if err := json.Unmarshal(message.Payload, &event); err != nil { - t.Fatalf("decode AOP payload: %v", err) - } - state, ok := sessions[event.SessionID] - if !ok { - continue - } - switch event.Type { - case aop.TypeMessage: - data, err := aop.DecodeData[aop.MessageData](event) - if err == nil && data.Role == "assistant" { - state.outputs = append(state.outputs, messageText(data)) - } - case aop.TypeTurnEnd: - state.ended++ - } - } - _ = cmd.Wait() - - for id, state := range sessions { - if state.ended != state.wantEnd { - t.Errorf("%s completed %d/%d runs (stderr: %s)", id, state.ended, state.wantEnd, clip(stderr.String(), 500)) - } - } - if got := strings.Join(sessions["sess-a"].outputs, "\n"); !strings.Contains(got, "ALPHA") { - t.Errorf("sess-a outputs = %q, want ALPHA", got) - } - if got := strings.Join(sessions["sess-b"].outputs, "\n"); !strings.Contains(got, "BRAVO") { - t.Errorf("sess-b outputs = %q, want BRAVO", got) - } -} - -func TestAgentEmptyReply(t *testing.T) { - h := New(t) - r := h.Agent("Reply with the word 'pong' and nothing else.") - Verify(t, r).OK().Done() - if !strings.Contains(strings.ToLower(r.Output()), "pong") { - t.Fatalf("expected 'pong', got: %s", r.Output()) - } -} - -func TestAgentBashTool(t *testing.T) { - h := New(t) - Intent{ - Name: "bash-echo", - Prompt: "Run 'echo hello_e2e' in a shell and tell me the exact output.", - Steps: Steps( - Tool("bash").ArgContains("echo hello_e2e").ResultHas("hello_e2e").NoError(), - ), - OutputContains: []string{"hello_e2e"}, - NoErrors: true, - MaxTurns: 3, - JudgeCriteria: "The agent must: (1) call the bash tool with a command containing 'echo hello_e2e', " + - "(2) the bash result must contain 'hello_e2e', " + - "(3) the final output must report 'hello_e2e' as the result.", - }.Run(t, h) -} - -func TestAgentReadTool(t *testing.T) { - h := New(t) - Intent{ - Name: "read-file", - Prompt: "Read /etc/hostname and reply with only its contents.", - Steps: Steps( - Tool("read").ArgContains("hostname").NoError(), - ), - NoErrors: true, - MaxTurns: 3, - JudgeCriteria: "The agent must use the read tool to read /etc/hostname, and the final output must contain the hostname value " + - "(not just say 'I read it' — the actual content must appear).", - }.Run(t, h) -} - -func TestAgentWriteReadRoundtrip(t *testing.T) { - h := New(t) - Intent{ - Name: "write-read-roundtrip", - Prompt: "Write 'e2e_marker_42' to /tmp/aiscan_e2e_test.txt, then read it back and confirm.", - Steps: Steps( - Tool("write").ArgContains("e2e_marker_42").NoError(), - Tool("read").ArgContains("aiscan_e2e_test").NoError(), - ), - Ordered: true, - OutputContains: []string{"e2e_marker_42"}, - NoErrors: true, - MaxTurns: 5, - JudgeCriteria: "The agent must: (1) write the exact string 'e2e_marker_42' to a file, " + - "(2) read it back and confirm the content matches. Both steps must succeed without errors.", - }.Run(t, h) -} - -func TestAgentGlobAndRead(t *testing.T) { - h := New(t) - Intent{ - Name: "glob-and-read", - Prompt: "List .go files in /mnt/chainreactors/aiscan/pkg/agent/ using glob, then read the first line of defaults.go and tell me the package name.", - Steps: Steps( - Tool("glob").NoError(), - Tool("read").ArgContains("defaults.go").NoError(), - ), - Ordered: true, - OutputContains: []string{"agent"}, - NoErrors: true, - MaxTurns: 4, - JudgeCriteria: "The agent must: (1) use glob to list .go files in the agent directory, " + - "(2) read defaults.go, (3) correctly report that the package name is 'agent'.", - }.Run(t, h) -} - -func TestAgentMultiStepTask(t *testing.T) { - h := New(t) - Intent{ - Name: "multi-step-bash", - Prompt: "First run 'uname -a' in bash. After you see the result, run 'whoami' in a SEPARATE bash call. Report both results.", - Steps: Steps( - Tool("bash").ArgContains("uname").NoError(), - Tool("bash").ArgContains("whoami").NoError(), - ), - Ordered: true, - NoErrors: true, - MaxTurns: 6, - JudgeCriteria: "The agent must make TWO separate bash calls: one for 'uname -a' and one for 'whoami'. " + - "Both results must appear in the final output. They must NOT be combined in a single bash call.", - }.Run(t, h) -} - -func TestAgentMultiTurn(t *testing.T) { - h := New(t) - Intent{ - Name: "multi-turn-file-ops", - Prompt: "Step 1: Create file /tmp/aiscan_multi.txt with content 'step1'. Step 2: Append ' step2' to it. Step 3: Read it and confirm it says 'step1 step2'.", - NoErrors: true, - MaxTurns: 8, - JudgeCriteria: "The agent must perform three sequential file operations: " + - "(1) create a file with 'step1', (2) append ' step2' to it, (3) read and confirm the content is 'step1 step2'. " + - "The final output must confirm the combined content.", - }.Run(t, h) -} - -func TestAgentLargeOutput(t *testing.T) { - h := New(t) - Intent{ - Name: "large-output", - Prompt: "Run 'seq 1 500' in bash. Tell me the last number printed.", - Steps: Steps( - Tool("bash").ArgContains("seq").NoError(), - ), - OutputContains: []string{"500"}, - NoErrors: true, - MaxTurns: 8, - JudgeCriteria: "The agent must run 'seq 1 500' and correctly identify that the last number is 500.", - }.Run(t, h) -} - -func TestAgentErrorRecovery(t *testing.T) { - h := New(t) - Intent{ - Name: "error-recovery", - Prompt: "Run 'cat /nonexistent/file' in bash. If it fails, report the error message. Then run 'echo recovered' and report that output.", - Steps: Steps( - Tool("bash").ArgContains("nonexistent"), - Tool("bash").ArgContains("recovered").NoError(), - ), - Ordered: true, - OutputContains: []string{"recovered"}, - MaxTurns: 5, - JudgeCriteria: "The agent must: (1) attempt to cat a nonexistent file, (2) recognize the error, " + - "(3) recover by running 'echo recovered', (4) report both the error and the recovery in the final output.", - }.Run(t, h) -} - -// ===================================================================== -// CLI — scanner help, version, direct modes -// ===================================================================== - -func TestScannerHelpExitsClean(t *testing.T) { - h := New(t) - for _, name := range scannerHelpCommands() { - t.Run(name, func(t *testing.T) { - r := h.Scanner(name, "-h") - Verify(t, r). - OK(). - OutputContains("Usage:"). - Done() - }) - } -} - -func TestVersionFlag(t *testing.T) { - h := New(t) - r := h.Run("--version") - Verify(t, r). - OK(). - OutputContains("aiscan v"). - Done() -} - -func TestScannerDirectGogo(t *testing.T) { - h := New(t) - r := h.Scanner("gogo", "-i", "127.0.0.1", "-p", "80") - if r.ExitCode != 0 { - t.Logf("gogo exit=%d stderr: %s", r.ExitCode, clip(r.Stderr, 500)) - } -} - -func TestScannerDirectSpray(t *testing.T) { - h := New(t) - r := h.Scanner("spray", "-i", "http://127.0.0.1:1", "--limit", "1") - if r.ExitCode != 0 { - t.Logf("spray exit=%d stderr: %s", r.ExitCode, clip(r.Stderr, 500)) - } -} - -func TestAgentTimeout(t *testing.T) { - h := New(t) - r := h.RunWithTimeout(15*time.Second, - "agent", "-p", "Run 'sleep 60' in bash.", - "--timeout", "5", - ) - if r.ExitCode == 0 && r.Duration < 4*time.Second { - t.Logf("agent completed before timeout — skipping assertion") - return - } - if r.Duration < 4*time.Second { - t.Fatalf("expected ≥4s duration, got %s", r.Duration) - } -} - -// ===================================================================== -// IOA loop — task dispatch, multi-worker, peer messages -// ===================================================================== - -func TestIOALoopReceivesTask(t *testing.T) { - service := ioaserver.NewService(ioaserver.NewMemoryStore(), "") - srv := httptest.NewServer(ioaserver.NewHandler(service)) - defer srv.Close() - - h := New(t) - - go func() { - h.RunWithTimeout(60*time.Second, - "agent", "--ioa-url", "http://127.0.0.1:8765", - "--ioa-url", srv.URL, - "--space", "test-loop", - "-p", "I am a test worker", - "--timeout", "45", - ) - }() - - time.Sleep(3 * time.Second) - - controller, err := ioaclient.NewClient(srv.URL, "") - if err != nil { - t.Fatal(err) - } - ctx := context.Background() - if _, err := controller.RegisterNode(ctx, "controller", "", nil); err != nil { - t.Fatal(err) - } - space, err := controller.Space(ctx, "test-loop", "e2e test") - if err != nil { - t.Fatal(err) - } - - nodes, err := controller.ListNodes(ctx) - if err != nil { - t.Fatal(err) - } - if len(nodes) == 0 { - t.Fatal("no worker nodes registered in space") - } - workerNodeID := nodes[0].ID - - _, err = controller.Send(ctx, space.ID, protocols.SendMessage{ - Content: map[string]any{"content": "Run 'echo ioa_task_received' in bash and report the output."}, - Refs: &protocols.Ref{Nodes: []string{workerNodeID}}, - }) - if err != nil { - t.Fatal(err) - } - - time.Sleep(30 * time.Second) - - requireIOAMessageContains(t, controller, ctx, space.ID, "ioa_task_received") -} - -func TestIOALoopMultipleWorkers(t *testing.T) { - service := ioaserver.NewService(ioaserver.NewMemoryStore(), "") - srv := httptest.NewServer(ioaserver.NewHandler(service)) - defer srv.Close() - - h := New(t) - - for i := 1; i <= 2; i++ { - i := i - go func() { - h.RunWithTimeout(45*time.Second, - "agent", "--ioa-url", "http://127.0.0.1:8765", - "--ioa-url", srv.URL, - "--space", "multi-worker", - "--ioa-node-name", fmt.Sprintf("worker-%d", i), - "-p", fmt.Sprintf("I am worker %d", i), - "--timeout", "40", - ) - }() - } - - time.Sleep(4 * time.Second) - - controller, err := ioaclient.NewClient(srv.URL, "") - if err != nil { - t.Fatal(err) - } - ctx := context.Background() - if _, err := controller.RegisterNode(ctx, "controller", "", nil); err != nil { - t.Fatal(err) - } - if _, err := controller.Space(ctx, "multi-worker", "e2e multi"); err != nil { - t.Fatal(err) - } - - nodes, err := controller.ListNodes(ctx) - if err != nil { - t.Fatal(err) - } - workerCount := 0 - for _, n := range nodes { - if strings.HasPrefix(n.Name, "worker-") { - workerCount++ - } - } - if workerCount < 2 { - t.Fatalf("expected ≥2 worker nodes, got %d (total nodes: %d)", workerCount, len(nodes)) - } -} - -func TestIOALoopPeerMessage(t *testing.T) { - service := ioaserver.NewService(ioaserver.NewMemoryStore(), "") - srv := httptest.NewServer(ioaserver.NewHandler(service)) - defer srv.Close() - - h := New(t) - - go func() { - h.RunWithTimeout(45*time.Second, - "agent", "--ioa-url", "http://127.0.0.1:8765", - "--ioa-url", srv.URL, - "--space", "peer-test", - "-p", "test worker", - "--timeout", "40", - ) - }() - - time.Sleep(3 * time.Second) - - controller, err := ioaclient.NewClient(srv.URL, "") - if err != nil { - t.Fatal(err) - } - ctx := context.Background() - if _, err := controller.RegisterNode(ctx, "controller", "", nil); err != nil { - t.Fatal(err) - } - space, err := controller.Space(ctx, "peer-test", "e2e peer") - if err != nil { - t.Fatal(err) - } - - nodes, err := controller.ListNodes(ctx) - if err != nil { - t.Fatal(err) - } - if len(nodes) == 0 { - t.Fatal("no worker nodes") - } - workerNodeID := nodes[0].ID - - _, err = controller.Send(ctx, space.ID, protocols.SendMessage{ - Content: map[string]any{"content": "Run echo peer_hello and report result"}, - Refs: &protocols.Ref{Nodes: []string{workerNodeID}}, - }) - if err != nil { - t.Fatal(err) - } - - _, err = controller.Send(ctx, space.ID, protocols.SendMessage{ - Content: map[string]any{"content": "Additional context: also run 'echo peer_context_received'"}, - }) - if err != nil { - t.Fatal(err) - } - - time.Sleep(25 * time.Second) - - requireIOAMessageContains(t, controller, ctx, space.ID, "peer_hello") -} - -func TestIOATaskSpawnsSubagents(t *testing.T) { - service := ioaserver.NewService(ioaserver.NewMemoryStore(), "") - srv := httptest.NewServer(ioaserver.NewHandler(service)) - defer srv.Close() - - h := New(t) - - go func() { - h.RunWithTimeout(90*time.Second, - "agent", "--ioa-url", "http://127.0.0.1:8765", - "--ioa-url", srv.URL, - "--space", "subagent-fan", - "-p", "I am a worker that parallelizes tasks using subagents", - "--timeout", "80", - ) - }() - - time.Sleep(4 * time.Second) - - controller, err := ioaclient.NewClient(srv.URL, "") - if err != nil { - t.Fatal(err) - } - ctx := context.Background() - if _, err := controller.RegisterNode(ctx, "controller", "", nil); err != nil { - t.Fatal(err) - } - space, err := controller.Space(ctx, "subagent-fan", "e2e") - if err != nil { - t.Fatal(err) - } - - nodes, err := controller.ListNodes(ctx) - if err != nil { - t.Fatal(err) - } - var workerNodeID string - for _, n := range nodes { - if n.Name != "controller" { - workerNodeID = n.ID - break - } - } - if workerNodeID == "" { - t.Fatal("no worker node found") - } - - _, err = controller.Send(ctx, space.ID, protocols.SendMessage{ - Content: map[string]any{ - "content": "I need you to gather system info in parallel. " + - "Create 2 async subagents: one runs 'echo subagent_alpha_ok' in bash, " + - "the other runs 'echo subagent_beta_ok' in bash. " + - "Wait for both results, then respond with a combined summary that includes both markers.", - }, - Refs: &protocols.Ref{Nodes: []string{workerNodeID}}, - }) - if err != nil { - t.Fatal(err) - } - - time.Sleep(60 * time.Second) - - requireIOAMessageContains(t, controller, ctx, space.ID, "subagent_alpha_ok") - requireIOAMessageContains(t, controller, ctx, space.ID, "subagent_beta_ok") -} - -func TestIOATwoWorkersDispatch(t *testing.T) { - service := ioaserver.NewService(ioaserver.NewMemoryStore(), "") - srv := httptest.NewServer(ioaserver.NewHandler(service)) - defer srv.Close() - - h := New(t) - - for i := 1; i <= 2; i++ { - i := i - go func() { - h.RunWithTimeout(75*time.Second, - "agent", "--ioa-url", "http://127.0.0.1:8765", - "--ioa-url", srv.URL, - "--space", "dispatch-2", - "--ioa-node-name", fmt.Sprintf("worker-%d", i), - "-p", fmt.Sprintf("I am worker %d", i), - "--timeout", "70", - ) - }() - } - - time.Sleep(5 * time.Second) - - controller, err := ioaclient.NewClient(srv.URL, "") - if err != nil { - t.Fatal(err) - } - ctx := context.Background() - if _, err := controller.RegisterNode(ctx, "controller", "", nil); err != nil { - t.Fatal(err) - } - space, err := controller.Space(ctx, "dispatch-2", "e2e dispatch") - if err != nil { - t.Fatal(err) - } - - nodes, err := controller.ListNodes(ctx) - if err != nil { - t.Fatal(err) - } - var workers []protocols.Node - for _, n := range nodes { - if strings.HasPrefix(n.Name, "worker-") { - workers = append(workers, n) - } - } - if len(workers) < 2 { - t.Fatalf("expected ≥2 workers, got %d", len(workers)) - } - - for i, w := range workers { - marker := fmt.Sprintf("dispatch_marker_%d", i+1) - _, err = controller.Send(ctx, space.ID, protocols.SendMessage{ - Content: map[string]any{ - "content": fmt.Sprintf("Run 'echo %s' in bash and report.", marker), - }, - Refs: &protocols.Ref{Nodes: []string{w.ID}}, - }) - if err != nil { - t.Fatal(err) - } - } - - time.Sleep(45 * time.Second) - - requireIOAMessageContains(t, controller, ctx, space.ID, "dispatch_marker_1") - requireIOAMessageContains(t, controller, ctx, space.ID, "dispatch_marker_2") -} - -// ===================================================================== -// Loop tool — create, lifecycle -// ===================================================================== - -func TestAgentLoopCreate(t *testing.T) { - h := New(t) - Intent{ - Name: "loop-create", - Prompt: "Use bash to run these loop commands in order: " + - "(1) loop '*/10 * * * *' check system health " + - "(2) loop list " + - "(3) loop stop the loop that was just created. " + - "Report the results and stop.", - Steps: Steps( - Tool("bash").ArgContains("loop").NoError(), - Tool("bash").ArgContains("loop").ArgContains("list").NoError(), - Tool("bash").ArgContains("loop").ArgContains("stop").NoError(), - ), - Ordered: true, - NoErrors: true, - MaxTurns: 6, - Timeout: 60 * time.Second, - JudgeCriteria: "The agent must: (1) create a loop via cron expression, " + - "(2) list loops, (3) stop the loop. All calls must succeed.", - }.Run(t, h) -} - -func TestAgentLoopLifecycle(t *testing.T) { - h := New(t) - Intent{ - Name: "loop-lifecycle", - Prompt: "Use bash to run these loop commands in order: " + - "(1) loop 5m check status " + - "(2) loop list to confirm the loop exists " + - "(3) loop stop to stop it " + - "(4) loop list again to confirm it is gone. " + - "Report the results after each step and stop.", - Steps: Steps( - Tool("bash").ArgContains("loop").NoError(), - Tool("bash").ArgContains("loop list").NoError(), - Tool("bash").ArgContains("loop stop").NoError(), - Tool("bash").ArgContains("loop list").NoError(), - ), - Ordered: true, - NoErrors: true, - MaxTurns: 8, - Timeout: 90 * time.Second, - JudgeCriteria: "The agent must: (1) create a loop, (2) list loops showing it exists, " + - "(3) stop the loop, (4) list loops again confirming it is gone. " + - "All four commands must succeed without errors.", - }.Run(t, h) -} - -// ===================================================================== -// Pipeline / scan — scanner AI, scan with skills -// ===================================================================== - -func TestScannerAIGogo(t *testing.T) { - h := New(t) - r := h.RunWithTimeout(90*time.Second, "--ai", "--timeout", "60", "gogo", "-i", "127.0.0.1", "-p", "80") - Verify(t, r).OK().Done() -} - -func TestAgentGogoScan(t *testing.T) { - h := New(t) - r := h.Agent("Use gogo to scan 127.0.0.1 port 80. Show the raw scanner output.") - Verify(t, r). - OK(). - ToolUsed("bash"). - Done() -} - -func TestAgentSprayScan(t *testing.T) { - h := New(t) - r := h.Agent("Run spray against http://127.0.0.1:1 with --limit 1 and report the result.") - Verify(t, r). - OK(). - ToolUsed("bash"). - Done() -} - -func TestAgentScanWithSkill(t *testing.T) { - h := New(t) - r := h.Agent("Use the scan command to scan 127.0.0.1 with --mode quick. Summarize the results.", "-s", "aiscan") - Verify(t, r).OK().Done() -} - -func TestAgentScanAnalyze(t *testing.T) { - h := New(t) - r := h.Agent("Run 'scan -i 127.0.0.1 --mode quick' and analyze the output. Tell me what services were found, if any.") - Verify(t, r). - OK(). - ToolUsed("bash"). - Done() -} - -func TestAgentScanAndVerify(t *testing.T) { - h := New(t) - r := h.Agent( - "Scan 127.0.0.1 with scan --mode quick. If any services are found, " + - "attempt to verify them by connecting to the reported port using bash (e.g. curl or nc). " + - "Report: services found, verification results.", - ) - Verify(t, r). - OK(). - ToolUsed("bash"). - Done() -} - -func TestAgentScanAnalyzeVerifyPipeline(t *testing.T) { - h := New(t) - r := h.Agent( - "Execute this pipeline:\n" + - "1. Run 'scan -i 127.0.0.1 --mode quick' to scan the target.\n" + - "2. Parse the scan results to identify any open ports or services.\n" + - "3. For each service found, attempt a basic verification:\n" + - " - If HTTP: run 'curl -s -o /dev/null -w \"%{http_code}\" http://127.0.0.1:' \n" + - " - If SSH: run 'echo | nc -w2 127.0.0.1 ' \n" + - " - If no services found, report that.\n" + - "4. Summarize: services found, verification status for each.", - ) - Verify(t, r). - OK(). - ToolUsed("bash"). - ToolArgMatch("bash", func(args string) bool { - return strings.Contains(args, "scan") && strings.Contains(args, "127.0.0.1") - }). - ToolResultMatch("bash", func(res string) bool { return res != "" }). - Done() -} - -func TestAgentParallelTargetScan(t *testing.T) { - h := New(t) - r := h.Agent( - "I need to check 3 targets in parallel. Create 3 async subagents:\n" + - "1. Named 'target-a': run 'echo target_a_scanned' in bash and report.\n" + - "2. Named 'target-b': run 'echo target_b_scanned' in bash and report.\n" + - "3. Named 'target-c': run 'echo target_c_scanned' in bash and report.\n" + - "Wait for ALL subagents to complete. List the subagents to track progress. " + - "Once all are done, produce a consolidated report with all 3 markers.", - ) - Verify(t, r). - OK(). - MinSubagentCreates(3). - OutputContains("target_a_scanned"). - OutputContains("target_b_scanned"). - OutputContains("target_c_scanned"). - Done() -} - -func TestAgentBackgroundTaskDrivesFollowUp(t *testing.T) { - h := New(t) - r := h.Agent( - "Start a detached tmux session: tmux new -d -s scan 'sleep 1 && echo SCAN_COMPLETE port=22 service=ssh'. " + - "Use tmux ls to confirm it's running. " + - "Use tmux wait -t scan to wait for it. Use tmux capture-pane -t scan to get output. " + - "Then run a follow-up command 'echo VERIFY_22_OK' to simulate verification. " + - "Report both the scan result and the verification result.", - ) - Verify(t, r). - OK(). - ToolUsed("bash"). - MinToolCalls(3). - AnyResultContains("SCAN_COMPLETE"). - AnyResultContains("VERIFY_22_OK"). - Done() -} - -func TestAgentTmuxAndSubagentCoordination(t *testing.T) { - h := New(t) - r := h.Agent( - "Do these in parallel:\n" + - "1. Start a detached tmux session: tmux new -d -s bg 'sleep 1 && echo bg_task_done_xyz'\n" + - "2. Create an async subagent named 'helper' with prompt: " + - "'Run echo subagent_helper_done in bash and report.'\n" + - "Monitor both: use tmux wait/capture-pane and wait for the subagent completion notification. " + - "Report both results when they complete.", - ) - Verify(t, r). - OK(). - ToolUsed("bash"). - ToolUsed("subagent"). - AnyResultContains("bg_task_done_xyz"). - AnyResultContains("subagent_helper_done"). - Done() -} - -// ===================================================================== -// Real scan — direct, AI, agent, subagent, loop, IOA -// ===================================================================== - -func sendMessage(content, nodeID string) protocols.SendMessage { - return protocols.SendMessage{ - Content: map[string]any{"content": content}, - Refs: &protocols.Ref{Nodes: []string{nodeID}}, - } -} - -const realTarget = "101.132.149.35/28" -const realSingleTarget = "101.132.149.35" - -// Layer 1: Direct scanner (no AI) — baseline - -func TestRealScanDirectGogo(t *testing.T) { - h := New(t) - r := h.RunWithTimeout(120*time.Second, "gogo", "-i", realTarget, "-p", "top100") - Verify(t, r).OK().Done() - t.Logf("gogo output (%d bytes):\n%s", len(r.Stdout), clip(r.Stdout, 2000)) -} - -func TestRealScanDirectSpray(t *testing.T) { - h := New(t) - r := h.RunWithTimeout(120*time.Second, "spray", "-i", fmt.Sprintf("http://%s", realSingleTarget), "--finger") - Verify(t, r).OK().Done() - t.Logf("spray output (%d bytes):\n%s", len(r.Stdout), clip(r.Stdout, 2000)) -} - -func TestRealScanDirectPipeline(t *testing.T) { - h := New(t) - r := h.RunWithTimeout(300*time.Second, "scan", "-i", realSingleTarget, "--mode", "quick") - Verify(t, r).OK().Done() - t.Logf("scan output (%d bytes):\n%s", len(r.Stdout), clip(r.Stdout, 3000)) -} - -// Layer 2: Scanner AI analysis and scan pipeline AI skills - -func TestRealScanGogoAI(t *testing.T) { - h := New(t) - Intent{ - Name: "real-gogo-ai", - Prompt: "", // not used in scanner AI mode - Timeout: 180 * time.Second, - JudgeCriteria: "The scanner must have executed gogo against the target and the AI must have provided " + - "a meaningful analysis of discovered services. The analysis should mention specific ports, " + - "services, or results - not just a generic summary.", - }.verifyScanner(t, h, "--ai", "--timeout", "120", "gogo", "-i", realTarget, "-p", "top100") -} - -func TestRealScanPipelineAISkills(t *testing.T) { - h := New(t) - Intent{ - Name: "real-scan-pipeline-ai-skills", - Prompt: "", - Timeout: 300 * time.Second, - JudgeCriteria: "The scan pipeline must have run against the target with explicit AI verification " + - "and sniper options. The output should include concrete scan findings or AI skill results, " + - "not just a generic completion message.", - }.verifyScanner(t, h, "--timeout", "240", "scan", "-i", realSingleTarget, "--mode", "quick", "--verify=high", "--sniper") -} - -// Layer 3: Agent mode — LLM decides how to scan - -func TestRealAgentGogoScan(t *testing.T) { - h := New(t) - Intent{ - Name: "real-agent-gogo", - Prompt: fmt.Sprintf("Use gogo to scan %s with port range top100. Report all discovered services including port, protocol, and any fingerprints.", realTarget), - Steps: Steps( - Tool("bash").ArgContains("gogo").NoError(), - ), - Timeout: 300 * time.Second, - MaxTurns: 20, - JudgeCriteria: "The agent must have executed gogo against 101.132.149.35/28 with appropriate port arguments. " + - "The final output must list specific discovered services (port numbers, service names). " + - "Generic statements like 'scan completed' without specific results are a failure.", - }.Run(t, h) -} - -func TestRealAgentSprayScan(t *testing.T) { - h := New(t) - Intent{ - Name: "real-agent-spray", - Prompt: fmt.Sprintf("Use spray to probe http://%s and identify web technologies and fingerprints. Report what you find.", realSingleTarget), - Steps: Steps( - Tool("bash").ArgContains("spray").NoError(), - ), - Timeout: 300 * time.Second, - MaxTurns: 20, - JudgeCriteria: "The agent must run spray against the target URL. The output must include specific web " + - "technology fingerprints or HTTP response information — not just 'spray completed'.", - }.Run(t, h) -} - -func TestRealAgentFullPipeline(t *testing.T) { - h := New(t) - Intent{ - Name: "real-agent-full-pipeline", - Prompt: fmt.Sprintf("Perform a comprehensive scan of %s:\n"+ - "1. Use gogo to discover open ports and services\n"+ - "2. For any HTTP services found, use spray to fingerprint them\n"+ - "3. Summarize all results: IPs, ports, services, web technologies", realSingleTarget), - Steps: Steps( - Tool("bash").ArgContains("gogo").NoError(), - ), - Timeout: 300 * time.Second, - MaxTurns: 12, - JudgeCriteria: "The agent must execute a multi-step scan: (1) port discovery with gogo, " + - "(2) web fingerprinting with spray for any HTTP services found. " + - "The final summary must list concrete results (specific IPs, ports, services). " + - "If no HTTP services are found, the agent should report that and skip spray — that's acceptable.", - }.Run(t, h) -} - -// Layer 4: Agent + skills - verify and analyze results - -func TestRealAgentScanWithVerify(t *testing.T) { - h := New(t) - Intent{ - Name: "real-agent-scan-verify", - Prompt: fmt.Sprintf("Scan %s with gogo. For each service found, attempt basic verification "+ - "(e.g. curl for HTTP, or nc for other services). Report: service, port, verification status.", realSingleTarget), - Steps: Steps( - Tool("bash").ArgContains("gogo").NoError(), - ), - Timeout: 300 * time.Second, - MaxTurns: 15, - JudgeCriteria: "The agent must: (1) run gogo to discover services, (2) attempt verification of at least one " + - "discovered service using curl/nc/similar. The report must show per-service verification status. " + - "If gogo finds no services, the agent should report that — still a pass if handled correctly.", - }.Run(t, h) -} - -func TestRealAgentScanReport(t *testing.T) { - h := New(t) - Intent{ - Name: "real-agent-scan-report", - Prompt: fmt.Sprintf("Scan %s using the scan command with --mode quick. Generate a security assessment report.", realSingleTarget), - Steps: Steps( - Tool("bash").ArgContains("scan").NoError(), - ), - Timeout: 300 * time.Second, - MaxTurns: 10, - JudgeCriteria: "The agent must run the scan pipeline and produce a structured security report. " + - "The report must contain: target IP, discovered services, risk assessment or observations. " + - "A bare scan output dump without analysis is a failure.", - }.Run(t, h) -} - -// Layer 5: Agent + subagent fan-out — parallel scanning - -func TestRealAgentParallelScan(t *testing.T) { - h := New(t) - Intent{ - Name: "real-agent-parallel-scan", - Prompt: fmt.Sprintf("I need to scan %s efficiently. Create 2 async subagents:\n"+ - "1. Named 'port-scan': run gogo against the target with -p top100\n"+ - "2. Named 'web-probe': run spray against http://%s with --finger\n"+ - "Wait for both to complete, then produce a consolidated results report.", realSingleTarget, realSingleTarget), - Steps: Steps( - Tool("subagent").Arg("name", "port-scan"), - Tool("subagent").Arg("name", "web-probe"), - ), - Timeout: 300 * time.Second, - MaxTurns: 12, - JudgeCriteria: "The agent must create 2 async subagents for parallel scanning. " + - "Both subagents must complete. The final report must consolidate results from both " + - "port scanning (gogo) and web probing (spray).", - }.Run(t, h) -} - -// Layer 6: Agent + loop tool — recurring scan - -func TestRealAgentLoopScan(t *testing.T) { - h := New(t) - Intent{ - Name: "real-agent-loop-scan", - Prompt: fmt.Sprintf("Set up a recurring scan for %s:\n"+ - "1. First, run gogo -i %s -p top100 immediately and report results\n"+ - "2. Create a loop named 'monitor' with interval '30s' and prompt 'check if any new ports opened on %s'\n"+ - "3. List loops to confirm the monitor is active\n"+ - "4. Delete the loop named 'monitor'\n"+ - "Report the initial scan results.", realSingleTarget, realSingleTarget, realSingleTarget), - Steps: Steps( - Tool("bash").ArgContains("gogo").NoError(), - ), - Ordered: true, - Timeout: 180 * time.Second, - MaxTurns: 10, - NoErrors: true, - JudgeCriteria: "The agent must: (1) run an initial gogo scan and report results, " + - "(2) create a recurring loop for monitoring, (3) list loops to confirm, (4) delete the loop. " + - "All four steps must complete in order. The initial scan must produce actual results (ports/services).", - }.Run(t, h) -} - -// Layer 7: IOA loop mode — swarm worker receives scan task - -func TestRealIOALoopScanTask(t *testing.T) { - service := ioaserver.NewService(ioaserver.NewMemoryStore(), "") - srv := httptest.NewServer(ioaserver.NewHandler(service)) - defer srv.Close() - - h := New(t) - - go func() { - h.RunWithTimeout(180*time.Second, - "agent", "--ioa-url", "http://127.0.0.1:8765", - "--ioa-url", srv.URL, - "--space", "real-scan", - "--ioa-node-name", "scanner-worker", - "-p", "I am a scanner worker with gogo, spray, and neutron capabilities", - "--timeout", "150", - ) - }() - - time.Sleep(5 * time.Second) - - controller, err := ioaclient.NewClient(srv.URL, "") - if err != nil { - t.Fatal(err) - } - ctx := context.Background() - if _, err := controller.RegisterNode(ctx, "controller", "", nil); err != nil { - t.Fatal(err) - } - space, err := controller.Space(ctx, "real-scan", "real scan test") - if err != nil { - t.Fatal(err) - } - - nodes, err := controller.ListNodes(ctx) - if err != nil { - t.Fatal(err) - } - var workerID string - for _, n := range nodes { - if n.Name == "scanner-worker" { - workerID = n.ID - break - } - } - if workerID == "" { - t.Fatal("scanner-worker not found") - } - - _, err = controller.Send(ctx, space.ID, sendMessage( - fmt.Sprintf("Run gogo against %s with -p top100 and report all discovered services with ports and fingerprints.", realSingleTarget), - workerID, - )) - if err != nil { - t.Fatal(err) - } - - time.Sleep(120 * time.Second) - - requireIOAMessageContains(t, controller, ctx, space.ID, realSingleTarget) -} - -// ===================================================================== -// Subagent — sync, async, fan-out, chain, message -// ===================================================================== - -func TestAgentSubagentSync(t *testing.T) { - h := New(t) - Intent{ - Name: "subagent-sync", - Prompt: "Use the subagent tool to create a sync subagent with prompt 'echo sub_sync_ok using bash and report the output'. Report the subagent result.", - Steps: Steps( - Tool("subagent").Action("create").NoError(), - ), - OutputContains: []string{"sub_sync_ok"}, - MaxTurns: 4, - JudgeCriteria: "The agent must create a sync subagent. The subagent must execute 'echo sub_sync_ok' via bash. " + - "The final output must contain 'sub_sync_ok' proving the subagent completed and returned its result.", - }.Run(t, h) -} - -func TestAgentSubagentAsync(t *testing.T) { - h := New(t) - Intent{ - Name: "subagent-async", - Prompt: "Create an async subagent with prompt 'Run echo async_marker_99 in bash'. Wait for its completion notification and report its result.", - Steps: Steps( - Tool("subagent").Action("create").NoError(), - ), - OutputContains: []string{"async_marker_99"}, - MaxTurns: 8, - JudgeCriteria: "The agent must create an async subagent. It must then wait for the subagent completion notification " + - "(which arrives via inbox). The final output must contain 'async_marker_99'.", - }.Run(t, h) -} - -func TestAgentSubagentSyncTimeout(t *testing.T) { - h := New(t) - Intent{ - Name: "subagent-sync-timeout", - Prompt: "Create a sync subagent with timeout '2s' and prompt 'Run sleep 30 in bash'. Report what happened (it should timeout).", - Steps: Steps( - Tool("subagent").ResultHas("timed out"), - ), - MaxTurns: 3, - JudgeCriteria: "The agent must create a sync subagent with a 2s timeout running 'sleep 30'. " + - "The subagent must timeout. The agent must report the timeout in its output.", - }.Run(t, h) -} - -func TestAgentSubagentList(t *testing.T) { - h := New(t) - Intent{ - Name: "subagent-list", - Prompt: "Create an async subagent named 'worker1' with prompt 'sleep 5'. Then immediately use subagent list action to show running subagents. Report the list.", - Steps: Steps( - Tool("subagent").Arg("name", "worker1"), - Tool("subagent").Action("list"), - ), - MaxTurns: 6, - JudgeCriteria: "The agent must: (1) create an async subagent named 'worker1', " + - "(2) call subagent list to show running subagents, " + - "(3) the list result should show 'worker1' as running.", - }.Run(t, h) -} - -func TestAgentMultiSubagentFanOut(t *testing.T) { - h := New(t) - Intent{ - Name: "subagent-fan-out", - Prompt: "You have 3 independent tasks. Use the subagent tool to create 3 SEPARATE async subagents, one for each:\n" + - "1. Subagent named 'host-info': run 'uname -a' in bash and report.\n" + - "2. Subagent named 'user-info': run 'whoami' in bash and report.\n" + - "3. Subagent named 'dir-info': run 'pwd' in bash and report.\n" + - "Create all 3 subagents, then wait for all completion notifications. " + - "Summarize all 3 results together.", - Steps: Steps( - Tool("subagent").Arg("name", "host-info"), - Tool("subagent").Arg("name", "user-info"), - Tool("subagent").Arg("name", "dir-info"), - ), - MaxTurns: 10, - JudgeCriteria: "The agent must create exactly 3 async subagents (host-info, user-info, dir-info). " + - "It must wait for all 3 completions. The final output must summarize results from all 3 subagents.", - }.Run(t, h) -} - -func TestAgentSubagentWithBashAndReport(t *testing.T) { - h := New(t) - Intent{ - Name: "subagent-bash-report", - Prompt: "Create 2 async subagents:\n" + - "1. Named 'counter': run 'seq 1 5' in bash.\n" + - "2. Named 'greeter': run 'echo hello_from_subagent' in bash.\n" + - "Wait for both to complete. Then report both outputs in your final answer.", - Steps: Steps( - Tool("subagent").Arg("name", "counter"), - Tool("subagent").Arg("name", "greeter"), - ), - OutputContains: []string{"hello_from_subagent"}, - MaxTurns: 10, - JudgeCriteria: "The agent must create 2 subagents and wait for both. " + - "The final output must include the output from both: the sequence 1-5 and 'hello_from_subagent'.", - }.Run(t, h) -} - -func TestAgentSubagentChain(t *testing.T) { - h := New(t) - Intent{ - Name: "subagent-chain", - Prompt: "Step 1: Create a sync subagent that runs 'echo chain_step_1' in bash and returns the output.\n" + - "Step 2: After you receive the result from step 1, create another sync subagent " + - "that runs 'echo chain_step_2' in bash.\n" + - "Report both results to confirm the chain completed.", - MaxTurns: 8, - JudgeCriteria: "The agent must create 2 sync subagents sequentially (not in parallel). " + - "Step 2 must happen AFTER step 1 completes. " + - "The final output must contain both 'chain_step_1' and 'chain_step_2'.", - Check: func(t *testing.T, r *RunResult) { - results := r.SubagentResults() - if len(results) < 2 { - t.Fatalf("expected ≥2 subagent results, got %d", len(results)) - } - s1, s2 := -1, -1 - for i, res := range results { - if strings.Contains(res, "chain_step_1") && s1 == -1 { - s1 = i - } - if strings.Contains(res, "chain_step_2") && s2 == -1 { - s2 = i - } - } - if s1 >= 0 && s2 >= 0 && s1 >= s2 { - t.Fatalf("chain order wrong: step1 at %d, step2 at %d", s1, s2) - } - }, - }.Run(t, h) -} - -func TestAgentSubagentMessage(t *testing.T) { - h := New(t) - Intent{ - Name: "subagent-message", - Prompt: "Create an async subagent named 'listener' with prompt: " + - "'Wait for a message. When you receive one, run echo GOT_MESSAGE in bash and report.'\n" + - "After creating it, use the subagent message action to send a message " + - "'hello from parent' to the 'listener' subagent.\n" + - "Wait for the listener to complete and report its result.", - Steps: Steps( - Tool("subagent").Arg("name", "listener"), - Tool("subagent").Action("message").Arg("name", "listener"), - ), - Ordered: true, - MaxTurns: 10, - JudgeCriteria: "The agent must: (1) create an async subagent named 'listener', " + - "(2) send a message to it via the subagent message action, " + - "(3) the listener must execute 'echo GOT_MESSAGE' after receiving the message, " + - "(4) the final output must contain 'GOT_MESSAGE' confirming the message was received and processed.", - }.Run(t, h) -} - -// ===================================================================== -// Task — tmux background tasks -// ===================================================================== - -func TestAgentBackgroundTask(t *testing.T) { - h := New(t) - r := h.Agent("Start a background shell session: tmux new -d -s bg 'sleep 1 && echo bg_done'. Then use tmux ls to list running sessions. Use tmux wait -t bg to wait for it to finish. Use tmux capture-pane -t bg to get the output. Report the final output.") - Verify(t, r). - OK(). - ToolUsed("bash"). - AnyResultContains("bg_done"). - NoToolErrors(). - Done() -} - -func TestAgentTmuxPeek(t *testing.T) { - h := New(t) - r := h.Agent("Run 'for i in 1 2 3; do echo line_$i; sleep 0.5; done' as a detached tmux session named 'lines'. Use tmux capture-pane -t lines --new to check its output, then wait for completion and report all lines.") - Verify(t, r). - OK(). - ToolUsed("bash"). - Done() -} - -func TestAgentTmuxKill(t *testing.T) { - h := New(t) - r := h.Agent("Start a detached tmux session: tmux new -d -s sleeper 'sleep 300'. Use tmux ls to confirm it's running. Kill it with tmux kill -t sleeper. List again to confirm it's killed. Report status.") - Verify(t, r). - OK(). - ToolUsed("bash"). - Done() -} - -// ===================================================================== -// Verify mechanism — scan verify/sniper mode tests -// ===================================================================== - -const verifyTarget = realSingleTarget - -// TestVerifyOffProducesNoAIOutput runs scan with --verify=off and confirms -// that no AI skill output appears in the results. -func TestVerifyOffProducesNoAIOutput(t *testing.T) { - h := New(t) - r := h.RunWithTimeout(300*time.Second, - "scan", "-i", "127.0.0.1", "--mode", "quick", "--verify=off", "--timeout", "3", - ) - Verify(t, r).OK().Done() - - if hasAISkillOutput(r.Stdout) { - t.Fatalf("--verify=off should produce no AI skill output, got:\n%s", clip(r.Stdout, 2000)) - } - t.Logf("verify=off output (%d bytes):\n%s", len(r.Stdout), clip(r.Stdout, 2000)) -} - -// TestVerifyHighWithSniperTriggersAIVerification runs scan with explicit -// verify and sniper options and confirms that the scan pipeline completes with -// AI skills enabled. When targets have high-priority loots, AI verify and -// sniper skills produce output. -func TestVerifyHighWithSniperTriggersAIVerification(t *testing.T) { - h := New(t) - r := h.RunWithTimeout(600*time.Second, - "scan", "-i", verifyTarget, "--mode", "quick", "--verify=high", "--sniper", "--timeout", "5", - ) - Verify(t, r).OK().Done() - - if !hasSummaryLine(r.Stdout) { - t.Fatal("expected [summary] line in output") - } - t.Logf("verify+sniper output (%d bytes):\n%s", len(r.Stdout), clip(r.Stdout, 3000)) -} - -// TestVerifyExplicitModeWithoutSniper runs scan with --verify=high explicitly -// (no --sniper) and checks that verify runs but sniper is NOT activated. -func TestVerifyExplicitModeWithoutSniper(t *testing.T) { - h := New(t) - r := h.RunWithTimeout(600*time.Second, - "scan", "-i", verifyTarget, "--mode", "quick", "--verify=high", "--timeout", "5", - ) - Verify(t, r).OK().Done() - - if hasSniperOutput(r.Stdout) { - t.Fatal("--verify=high without --sniper should not produce sniper output") - } - if !hasSummaryLine(r.Stdout) { - t.Fatal("expected [summary] line in output") - } - t.Logf("verify=high output (%d bytes):\n%s", len(r.Stdout), clip(r.Stdout, 2000)) -} - -// TestScanVerifySniperNoPostAnalysis verifies that the old post-analysis -// one-shot LLM call no longer runs. Explicit scan AI skills trigger only -// in-pipeline AI work (verify + sniper), not a separate "analysis" step. -// The output should contain the [summary] line from the scan pipeline but -// should not contain the "analysis" output section that runScannerPostAnalysis -// used to produce. -func TestScanVerifySniperNoPostAnalysis(t *testing.T) { - h := New(t) - r := h.RunWithTimeout(600*time.Second, - "scan", "-i", verifyTarget, "--mode", "quick", "--verify=high", "--sniper", "--timeout", "5", - ) - Verify(t, r).OK().Done() - - if !hasSummaryLine(r.Stdout) { - t.Fatal("expected [summary] line from scan pipeline") - } - t.Logf("output (%d bytes):\n%s", len(r.Stdout), clip(r.Stdout, 3000)) -} - -// TestScanDefaultModeCompletes runs scan without any explicit AI skill flags. -// The default verify mode is "auto" (mapped to "high"), which enables the -// provider optionally. If the provider initializes, AI verify can run; if not, -// the scan still completes successfully. -func TestScanDefaultModeCompletes(t *testing.T) { - h := New(t) - r := h.RunWithTimeout(600*time.Second, - "scan", "-i", verifyTarget, "--mode", "quick", "--timeout", "5", - ) - Verify(t, r).OK().Done() - - if !hasSummaryLine(r.Stdout) { - t.Fatal("expected [summary] line in output") - } - t.Logf("default mode output (%d bytes):\n%s", len(r.Stdout), clip(r.Stdout, 2000)) -} - -// TestVerifyOffDisablesAllAISkills confirms that --verify=off combined with -// no --sniper and no --deep results in zero AI skill results in the summary. -func TestVerifyOffDisablesAllAISkills(t *testing.T) { - h := New(t) - r := h.RunWithTimeout(300*time.Second, - "scan", "-i", "127.0.0.1", "--mode", "quick", "--verify=off", "--timeout", "3", - ) - Verify(t, r).OK().Done() - - summary := extractSummaryLine(r.Stdout) - if summary == "" { - t.Fatal("missing [summary] line") - } - if strings.Contains(summary, "verified") { - parts := strings.Fields(summary) - for i, p := range parts { - if p == "verified" && i > 0 && parts[i-1] != "0" { - t.Fatalf("expected 0 verified in summary with --verify=off, got: %s", summary) - } - } - } - t.Logf("verify=off summary: %s", summary) -} - -// TestScanVerifyWithReportIncludesVerification runs scan with explicit -// verification and report output and verifies the report includes AI -// verification metrics. -func TestScanVerifyWithReportIncludesVerification(t *testing.T) { - h := New(t) - r := h.RunWithTimeout(600*time.Second, - "scan", "-i", verifyTarget, "--mode", "quick", "--verify=high", "--sniper", "--report", "--timeout", "5", - ) - Verify(t, r).OK().Done() - - hasMetrics := strings.Contains(r.Stdout, "AI verifications") || - strings.Contains(r.Stdout, "AI skill") || - strings.Contains(r.Stdout, "verified") - if !hasMetrics { - t.Fatal("--verify=high --sniper --report should include AI verification information in output") - } - t.Logf("report output (%d bytes):\n%s", len(r.Stdout), clip(r.Stdout, 3000)) -} - -// TestAssetReportFileOutputFormats runs scan with -f and -F and verifies both -// output formats include structured checkpoint loots. -func TestAssetReportFileOutputFormats(t *testing.T) { - h := New(t) - r := h.RunWithTimeout(600*time.Second, - "scan", "-i", verifyTarget, "--mode", "quick", "--verify=high", "--sniper", "--timeout", "5", - "-f", "output.txt", "-F", "asset_report.txt", - ) - Verify(t, r).OK().Done() - - plainBytes, err := os.ReadFile(h.WorkFile("output.txt")) - if err != nil { - t.Fatalf("read -f output: %v", err) - } - plain := string(plainBytes) - t.Logf("-f output (%d bytes):\n%s", len(plain), clip(plain, 3000)) - - assetReportBytes, err := os.ReadFile(h.WorkFile("asset_report.txt")) - if err != nil { - if !hasAISkillOutput(r.Stdout) { - t.Skip("no AI output produced, skipping -F check") - } - t.Fatalf("read -F output: %v", err) - } - assetReport := string(assetReportBytes) - t.Logf("-F output (%d bytes):\n%s", len(assetReport), clip(assetReport, 3000)) - - if len(assetReport) > 0 { - if !strings.Contains(assetReport, "Assets:") { - t.Fatal("-F output should contain 'Assets:' header") - } - } -} - -// ===================================================================== -// Shared helpers -// ===================================================================== - -// containsCount counts occurrences of substr in s. -func containsCount(s, substr string) int { - return strings.Count(s, substr) -} - -// requireIOAMessageContains checks that at least one message in the space contains substr. -func requireIOAMessageContains(t *testing.T, client *ioaclient.Client, ctx context.Context, spaceID, substr string) { - t.Helper() - msgs, err := client.Read(ctx, spaceID, protocols.ReadOptions{All: true}) - if err != nil { - t.Fatalf("read space: %v", err) - } - for _, m := range msgs { - raw, _ := json.Marshal(m.Content) - if strings.Contains(string(raw), substr) { - return - } - } - var summaries []string - for _, m := range msgs { - raw, _ := json.Marshal(m.Content) - summaries = append(summaries, clip(string(raw), 200)) - } - t.Fatalf("no IOA message contains %q:\n%s", substr, strings.Join(summaries, "\n")) -} - -// verifyScanner runs a direct scanner command and uses the judge to evaluate. -func (intent Intent) verifyScanner(t *testing.T, h *Harness, args ...string) *RunResult { - t.Helper() - r := h.RunWithTimeout(intent.Timeout, args...) - v := Verify(t, r).OK() - if intent.JudgeCriteria != "" { - prompt := fmt.Sprintf("Scanner command: %v", args) - v = v.JudgeWith(h.Judge(), prompt, intent.JudgeCriteria) - } - v.Done() - t.Logf("output (%d bytes):\n%s", len(r.Stdout), clip(r.Stdout, 2000)) - return r -} - -func hasAISkillOutput(output string) bool { - markers := []string{"[ai:", "[sniper:", "[ai]", "[sniper]"} - for _, m := range markers { - if strings.Contains(output, m) { - return true - } - } - return false -} - -func hasSniperOutput(output string) bool { - return strings.Contains(output, "[sniper:") || strings.Contains(output, "[sniper]") -} - -func hasSummaryLine(output string) bool { - return strings.Contains(output, "[summary]") || strings.Contains(output, "completed") -} - -func extractSummaryLine(output string) string { - for _, line := range strings.Split(output, "\n") { - if strings.Contains(line, "[summary]") { - return line - } - } - return "" -} diff --git a/core/harness/intent.go b/core/harness/intent.go deleted file mode 100644 index e5b57d2a..00000000 --- a/core/harness/intent.go +++ /dev/null @@ -1,167 +0,0 @@ -//go:build e2e - -package harness - -import ( - "fmt" - "strings" - "testing" - "time" -) - -// Intent describes a complete AI behavior test case declaratively. -// Instead of writing imperative test code, define an Intent and call Run. -// -// Intent{ -// Name: "subagent-lifecycle", -// Prompt: "Create a sync subagent to scan localhost.", -// Steps: Steps( -// Tool("subagent").Action("create").Arg("name", "scanner"), -// ), -// Ordered: true, -// MaxTurns: 4, -// NoErrors: true, -// }.Run(t, h) -type Intent struct { - Name string - Prompt string - ExtraArgs []string - Timeout time.Duration - - // Steps describes expected tool calls. - Steps []ToolPattern - - // Ordered requires steps to appear in sequence (subsequence match). - // When false, steps can appear in any order. - Ordered bool - - // OutputContains lists substrings that must appear in stdout/stderr. - OutputContains []string - - // OutputMissing lists substrings that must NOT appear in output. - OutputMissing []string - - // NoErrors requires all tool calls to succeed. - NoErrors bool - - // MaxTurns caps the number of turns (0 = no limit). - MaxTurns int - - // MaxToolCalls caps total tool invocations (0 = no limit). - MaxToolCalls int - - // MaxDuration caps wall-clock time (0 = no limit). - MaxDuration time.Duration - - // JudgeCriteria, when non-empty, enables LLM-as-judge evaluation. - // The judge receives the intent prompt, this criteria string, and the - // full execution trace. It returns a pass/fail verdict. - // Example: "The agent must have created exactly one loop named 'scanner', - // listed it to confirm it exists, then deleted it." - JudgeCriteria string - - // Check is an optional custom verification function. - Check func(t *testing.T, r *RunResult) -} - -// Steps is a convenience constructor for []ToolPattern. -func Steps(patterns ...ToolPattern) []ToolPattern { return patterns } - -// Run executes the intent against the harness and verifies all expectations. -func (intent Intent) Run(t *testing.T, h *Harness) *RunResult { - t.Helper() - - var r *RunResult - if intent.Timeout > 0 { - r = h.AgentWithTimeout(intent.Timeout, intent.Prompt, intent.ExtraArgs...) - } else { - r = h.Agent(intent.Prompt, intent.ExtraArgs...) - } - intent.verify(t, h, r) - return r -} - -func (intent Intent) verify(t *testing.T, h *Harness, r *RunResult) { - t.Helper() - - v := Verify(t, r).OK() - - // structural checks - if len(intent.Steps) > 0 { - if intent.Ordered { - v = v.ExpectInOrder(intent.Steps...) - } else { - v = v.Expect(intent.Steps...) - } - } - for _, s := range intent.OutputContains { - v = v.OutputContains(s) - } - for _, s := range intent.OutputMissing { - v = v.OutputMissing(s) - } - if intent.NoErrors { - v = v.NoToolErrors() - } - if intent.MaxTurns > 0 { - v = v.MaxTurns(intent.MaxTurns) - } - if intent.MaxToolCalls > 0 { - v = v.MaxToolCalls(intent.MaxToolCalls) - } - if intent.MaxDuration > 0 { - v = v.CompletedWithin(intent.MaxDuration) - } - - // semantic check via LLM judge - if intent.JudgeCriteria != "" { - v = v.JudgeWith(h.Judge(), intent.Prompt, intent.JudgeCriteria) - } - - v.Done() - - if intent.Check != nil { - intent.Check(t, r) - } -} - -// Describe returns a human-readable summary of the intent for logging. -func (intent Intent) Describe() string { - var sb strings.Builder - sb.WriteString(fmt.Sprintf("Intent: %s\n", intent.Name)) - sb.WriteString(fmt.Sprintf(" Prompt: %s\n", clip(intent.Prompt, 80))) - if len(intent.Steps) > 0 { - order := "any order" - if intent.Ordered { - order = "in order" - } - sb.WriteString(fmt.Sprintf(" Steps (%s):\n", order)) - for i, s := range intent.Steps { - sb.WriteString(fmt.Sprintf(" %d. %s\n", i+1, s.describe())) - } - } - if len(intent.OutputContains) > 0 { - sb.WriteString(fmt.Sprintf(" Output must contain: %v\n", intent.OutputContains)) - } - if intent.MaxTurns > 0 { - sb.WriteString(fmt.Sprintf(" Max turns: %d\n", intent.MaxTurns)) - } - if intent.NoErrors { - sb.WriteString(" No tool errors allowed\n") - } - return sb.String() -} - -// IntentSuite runs multiple intents as subtests. -func IntentSuite(t *testing.T, h *Harness, intents ...Intent) { - t.Helper() - for _, intent := range intents { - name := intent.Name - if name == "" { - name = clip(intent.Prompt, 40) - } - t.Run(name, func(t *testing.T) { - intent.Run(t, h) - }) - } -} diff --git a/core/harness/judge.go b/core/harness/judge.go deleted file mode 100644 index 7fe26a93..00000000 --- a/core/harness/judge.go +++ /dev/null @@ -1,205 +0,0 @@ -//go:build e2e - -package harness - -import ( - "bytes" - "context" - "encoding/json" - "fmt" - "io" - "net/http" - "strings" - "time" -) - -// Verdict is the structured result from an LLM judge evaluation. -type Verdict struct { - Pass bool `json:"pass"` - Score int `json:"score"` - Reason string `json:"reason"` - Issues []string `json:"issues"` -} - -// Judge evaluates agent execution results using an LLM. -type Judge struct { - baseURL string - apiKey string - model string - timeout time.Duration -} - -func NewJudge(baseURL, apiKey, model string) *Judge { - return &Judge{ - baseURL: strings.TrimRight(baseURL, "/"), - apiKey: apiKey, - model: model, - timeout: 30 * time.Second, - } -} - -func (h *Harness) Judge() *Judge { - return NewJudge(h.baseURL, h.apiKey, h.model) -} - -const judgeMaxRetries = 3 - -// Evaluate sends the intent and execution trace to the LLM for judgment. -func (j *Judge) Evaluate(intent string, criteria string, r *RunResult) (*Verdict, error) { - trace := buildTrace(r) - prompt := buildJudgePrompt(intent, criteria, trace) - - var lastErr error - for attempt := 0; attempt < judgeMaxRetries; attempt++ { - v, err := j.call(prompt) - if err == nil { - return v, nil - } - lastErr = err - if attempt < judgeMaxRetries-1 { - time.Sleep(time.Duration(attempt+1) * time.Second) - } - } - return nil, fmt.Errorf("judge failed after %d attempts: %w", judgeMaxRetries, lastErr) -} - -func buildTrace(r *RunResult) string { - var sb strings.Builder - fmt.Fprintf(&sb, "Exit code: %d\n", r.ExitCode) - fmt.Fprintf(&sb, "Duration: %s\n", r.Duration.Round(time.Millisecond)) - fmt.Fprintf(&sb, "Turns: %d\n", r.Turns()) - fmt.Fprintf(&sb, "Tool calls: %d\n", len(r.ToolCalls())) - - sb.WriteString("\nTool call trace:\n") - for i, e := range r.ToolCalls() { - fmt.Fprintf(&sb, " [%d] %s", i+1, e.Name()) - if e.IsError() { - sb.WriteString(" (ERROR)") - } - sb.WriteByte('\n') - if args := argsText(e.Args()); args != "" { - fmt.Fprintf(&sb, " args: %s\n", clip(args, 200)) - } - if result := e.ResultText(); result != "" { - fmt.Fprintf(&sb, " result: %s\n", clip(result, 300)) - } - } - - if output := strings.TrimSpace(r.Stdout); output != "" { - fmt.Fprintf(&sb, "\nFinal output:\n%s\n", clip(output, 1000)) - } - return sb.String() -} - -const judgeSystemPrompt = `You are a strict test evaluator for an AI agent system. Given an intent (what was asked), evaluation criteria, and execution trace (what happened), determine whether the agent correctly fulfilled the intent. - -Respond with ONLY a JSON object: -{"pass": true/false, "score": 0-100, "reason": "one sentence summary", "issues": ["issue1", "issue2"]} - -Rules: -- pass=true only if the intent was fully and correctly completed -- score: 100=perfect, 80+=good, 60+=acceptable, <60=fail -- issues: list specific problems (empty if pass=true) -- Be strict: "ran without errors" is not the same as "fulfilled the intent" -- Check that the right tools were used with correct arguments -- Check that results contain expected data, not just that tools were called` - -func buildJudgePrompt(intent, criteria, trace string) string { - var sb strings.Builder - fmt.Fprintf(&sb, "## Intent\n%s\n\n", intent) - if criteria != "" { - fmt.Fprintf(&sb, "## Evaluation Criteria\n%s\n\n", criteria) - } - fmt.Fprintf(&sb, "## Execution Trace\n%s", trace) - return sb.String() -} - -type chatRequest struct { - Model string `json:"model"` - Messages []chatMessage `json:"messages"` - MaxTokens int `json:"max_tokens"` - Temperature float64 `json:"temperature"` -} - -type chatMessage struct { - Role string `json:"role"` - Content string `json:"content"` -} - -type chatResponse struct { - Choices []struct { - Message struct { - Content string `json:"content"` - } `json:"message"` - } `json:"choices"` -} - -func (j *Judge) call(userPrompt string) (*Verdict, error) { - body := chatRequest{ - Model: j.model, - Messages: []chatMessage{ - {Role: "system", Content: judgeSystemPrompt}, - {Role: "user", Content: userPrompt}, - }, - MaxTokens: 512, - Temperature: 0, - } - - data, err := json.Marshal(body) - if err != nil { - return nil, fmt.Errorf("marshal request: %w", err) - } - - ctx, cancel := context.WithTimeout(context.Background(), j.timeout) - defer cancel() - - url := j.baseURL + "/chat/completions" - req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(data)) - if err != nil { - return nil, err - } - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", "Bearer "+j.apiKey) - - resp, err := http.DefaultClient.Do(req) - if err != nil { - return nil, fmt.Errorf("judge API call failed: %w", err) - } - defer resp.Body.Close() - - respData, err := io.ReadAll(resp.Body) - if err != nil { - return nil, fmt.Errorf("read response: %w", err) - } - if resp.StatusCode != 200 { - return nil, fmt.Errorf("judge API returned %d: %s", resp.StatusCode, clip(string(respData), 500)) - } - - var chatResp chatResponse - if err := json.Unmarshal(respData, &chatResp); err != nil { - return nil, fmt.Errorf("parse response: %w", err) - } - if len(chatResp.Choices) == 0 { - return nil, fmt.Errorf("judge returned no choices") - } - - return parseVerdict(chatResp.Choices[0].Message.Content) -} - -func parseVerdict(raw string) (*Verdict, error) { - raw = strings.TrimSpace(raw) - raw = stripJSONFences(raw) - - var v Verdict - if err := json.Unmarshal([]byte(raw), &v); err != nil { - return nil, fmt.Errorf("parse verdict JSON: %w\nraw: %s", err, clip(raw, 500)) - } - return &v, nil -} - -func stripJSONFences(s string) string { - s = strings.TrimPrefix(s, "```json") - s = strings.TrimPrefix(s, "```") - s = strings.TrimSuffix(s, "```") - return strings.TrimSpace(s) -} diff --git a/core/harness/monitor.go b/core/harness/monitor.go deleted file mode 100644 index 9b7da9dc..00000000 --- a/core/harness/monitor.go +++ /dev/null @@ -1,76 +0,0 @@ -//go:build e2e - -package harness - -import ( - "fmt" - "io" - - "github.com/chainreactors/aiscan/pkg/agent/truncate" - "github.com/chainreactors/aiscan/pkg/aop" -) - -// Monitor renders the AOP events received from the agent's webproto stdout. -// Attach it to a Harness with h.WithMonitor(). -// -// Output goes to the provided Writer (typically os.Stderr for live -// terminal view, or a test log adapter). -type Monitor struct { - out io.Writer - runSeen string -} - -func NewMonitor(out io.Writer) *Monitor { - return &Monitor{out: out} -} - -func (m *Monitor) printf(format string, args ...any) { - fmt.Fprintf(m.out, format, args...) -} - -func (m *Monitor) renderEvent(ev aop.Event) { - switch ev.Type { - case aop.TypeSessionStart: - m.runSeen = "" - - case aop.TypeTurnStart: - if ev.TurnID != "" && ev.TurnID != m.runSeen { - m.runSeen = ev.TurnID - m.printf("\n── run %s ──\n", ev.TurnID) - } - - case aop.TypeMessage: - data, err := aop.DecodeData[aop.MessageData](ev) - if err == nil && (data.Role == "" || data.Role == "assistant") { - if text := messageText(data); text != "" { - m.printf(" 💬 %s\n", truncate.Clip(text, 200)) - } - } - - case aop.TypeToolCall: - data, err := aop.DecodeData[aop.ToolCallData](ev) - if err == nil { - m.printf(" 🔧 %s %s\n", data.ToolName, truncate.Clip(argsText(data.Args), 120)) - } - - case aop.TypeToolResult: - data, err := aop.DecodeData[aop.ToolResultData](ev) - if err == nil { - result := valueText(data.Content) - switch { - case data.IsError: - m.printf(" ❌ %s error: %s\n", data.ToolName, truncate.Clip(result, 100)) - case result != "": - m.printf(" ✓ %s → %d bytes: %s\n", data.ToolName, len(result), truncate.Clip(result, 100)) - default: - m.printf(" ✓ %s → (empty)\n", data.ToolName) - } - } - - case aop.TypeTurnEnd: - data, err := aop.DecodeData[aop.TurnEndData](ev) - if err == nil { - m.printf("\n── run done (stop=%s) ──\n", data.Stop) - } - } -} diff --git a/core/harness/result.go b/core/harness/result.go deleted file mode 100644 index 6918cb8f..00000000 --- a/core/harness/result.go +++ /dev/null @@ -1,233 +0,0 @@ -//go:build e2e - -package harness - -import ( - "encoding/json" - "strings" - "time" - - "github.com/chainreactors/aiscan/pkg/aop" -) - -type RunResult struct { - Stdout string - Stderr string - ExitCode int - Duration time.Duration - Events []aop.Event -} - -// ToolExecution is an on-demand typed view that retains both original AOP -// envelopes. It is never stored in place of the protocol events. -type ToolExecution struct { - CallEvent aop.Event - ResultEvent aop.Event - Call aop.ToolCallData - Result aop.ToolResultData -} - -func (e ToolExecution) Name() string { - if e.Result.ToolName != "" { - return e.Result.ToolName - } - return e.Call.ToolName -} - -func (e ToolExecution) Args() any { return e.Call.Args } -func (e ToolExecution) ResultText() string { return valueText(e.Result.Content) } -func (e ToolExecution) IsError() bool { return e.Result.IsError } - -func (r *RunResult) OK() bool { return r.ExitCode == 0 } -func (r *RunResult) Output() string { return strings.TrimSpace(r.Stdout) } -func (r *RunResult) Combined() string { return r.Stdout + r.Stderr } - -func (r *RunResult) ContainsOutput(substr string) bool { - return strings.Contains(r.Stdout, substr) || strings.Contains(r.Stderr, substr) -} - -func (r *RunResult) ToolCalls() []ToolExecution { - calls := make(map[string]struct { - event aop.Event - data aop.ToolCallData - }) - for _, event := range r.Events { - if event.Type != aop.TypeToolCall { - continue - } - data, err := aop.DecodeData[aop.ToolCallData](event) - if err == nil && data.ToolCallID != "" { - calls[data.ToolCallID] = struct { - event aop.Event - data aop.ToolCallData - }{event: event, data: data} - } - } - var out []ToolExecution - for _, event := range r.Events { - if event.Type != aop.TypeToolResult { - continue - } - data, err := aop.DecodeData[aop.ToolResultData](event) - if err != nil { - continue - } - call := calls[data.ToolCallID] - out = append(out, ToolExecution{ - CallEvent: call.event, ResultEvent: event, Call: call.data, Result: data, - }) - } - return out -} - -func (r *RunResult) HasToolCall(name string) bool { - return len(r.ToolCallsNamed(name)) > 0 -} - -func (r *RunResult) ToolCallsNamed(name string) []ToolExecution { - var out []ToolExecution - for _, execution := range r.ToolCalls() { - if execution.Name() == name { - out = append(out, execution) - } - } - return out -} - -func (r *RunResult) Turns() int { - seen := make(map[string]struct{}) - for _, event := range r.Events { - if event.Type == aop.TypeTurnStart && event.TurnID != "" { - seen[event.TurnID] = struct{}{} - } - } - return len(seen) -} - -func (r *RunResult) ToolCallSequence() []string { - var names []string - for _, execution := range r.ToolCalls() { - names = append(names, execution.Name()) - } - return names -} - -func (r *RunResult) ToolResultContains(toolName, substr string) bool { - for _, execution := range r.ToolCallsNamed(toolName) { - if strings.Contains(execution.ResultText(), substr) { - return true - } - } - return false -} - -func (r *RunResult) ToolArgsContains(toolName, substr string) bool { - for _, execution := range r.ToolCallsNamed(toolName) { - if strings.Contains(argsText(execution.Args()), substr) { - return true - } - } - return false -} - -func (r *RunResult) AllToolResults() string { - var sb strings.Builder - for _, execution := range r.ToolCalls() { - sb.WriteString(execution.ResultText()) - sb.WriteByte('\n') - } - return sb.String() -} - -func (r *RunResult) ErroredToolCalls() []ToolExecution { - var out []ToolExecution - for _, execution := range r.ToolCalls() { - if execution.IsError() { - out = append(out, execution) - } - } - return out -} - -func (r *RunResult) StopReason() string { - for i := len(r.Events) - 1; i >= 0; i-- { - if r.Events[i].Type != aop.TypeTurnEnd { - continue - } - data, err := aop.DecodeData[aop.TurnEndData](r.Events[i]) - if err == nil { - return data.Stop - } - } - return "" -} - -func (r *RunResult) TotalTokens() int { - for i := len(r.Events) - 1; i >= 0; i-- { - if r.Events[i].Type != aop.TypeUsage { - continue - } - data, err := aop.DecodeData[aop.UsageData](r.Events[i]) - if err == nil && data.TotalTokens > 0 { - return data.TotalTokens - } - } - return 0 -} - -func (r *RunResult) SubagentCalls() []ToolExecution { return r.ToolCallsNamed("subagent") } - -func (r *RunResult) SubagentCreateCount() int { - n := 0 - for _, execution := range r.SubagentCalls() { - if isSubagentCreate(execution) { - n++ - } - } - return n -} - -func (r *RunResult) SubagentCreateArgs() []string { - var args []string - for _, execution := range r.SubagentCalls() { - if isSubagentCreate(execution) { - args = append(args, argsText(execution.Args())) - } - } - return args -} - -func (r *RunResult) SubagentResults() []string { - var results []string - for _, execution := range r.SubagentCalls() { - if isSubagentCreate(execution) { - results = append(results, execution.ResultText()) - } - } - return results -} - -func isSubagentCreate(execution ToolExecution) bool { - args, ok := execution.Args().(map[string]any) - if !ok { - return true - } - action, _ := args["action"].(string) - return action != "list" && action != "kill" && action != "message" -} - -func argsText(args any) string { return valueText(args) } - -func valueText(value any) string { - if value == nil { - return "" - } - if text, ok := value.(string); ok { - return text - } - encoded, err := json.Marshal(value) - if err != nil { - return "" - } - return string(encoded) -} diff --git a/core/harness/stdio.go b/core/harness/stdio.go deleted file mode 100644 index ac74aa59..00000000 --- a/core/harness/stdio.go +++ /dev/null @@ -1,128 +0,0 @@ -//go:build e2e - -package harness - -import ( - "encoding/json" - "errors" - "fmt" - "io" - "strings" - - "github.com/chainreactors/aiscan/pkg/aop" - "github.com/chainreactors/aiscan/pkg/webproto" -) - -// messageText flattens the text parts of a complete AOP message. -func messageText(data aop.MessageData) string { - var sb strings.Builder - for _, part := range data.Parts { - if part.Type != aop.PartText || part.Text == "" { - continue - } - if sb.Len() > 0 { - sb.WriteString("\n") - } - sb.WriteString(part.Text) - } - return sb.String() -} - -func consumeAgentStream(input io.Reader, monitor *Monitor) (string, []aop.Event, error) { - var output string - var events []aop.Event - decoder := json.NewDecoder(input) - sessionID := "" - turnID := "" - var agentErr error - - for { - var message webproto.Message - if err := decoder.Decode(&message); err != nil { - if errors.Is(err, io.EOF) { - if sessionID == "" { - return output, events, fmt.Errorf("stdio stream ended without session.opened") - } - return output, events, fmt.Errorf("stdio stream ended without run turn.end") - } - return output, events, fmt.Errorf("decode stdio frame: %w", err) - } - switch message.Type { - case webproto.TypeSessionOpened: - var data webproto.SessionLifecyclePayload - if err := json.Unmarshal(message.Payload, &data); err != nil { - return output, events, fmt.Errorf("decode session.opened: %w", err) - } - if data.SessionID == "" { - return output, events, fmt.Errorf("session.opened has empty session_id") - } - if sessionID == "" { - sessionID = data.SessionID - } else if sessionID != data.SessionID { - return output, events, fmt.Errorf("unexpected session.opened for %q", data.SessionID) - } - - case webproto.TypeError: - var data webproto.ErrorPayload - if err := json.Unmarshal(message.Payload, &data); err != nil { - return output, events, fmt.Errorf("decode error frame: %w", err) - } - if strings.TrimSpace(data.Message) == "" { - return output, events, fmt.Errorf("stdio error frame has empty message") - } - return output, events, fmt.Errorf("agent error: %s", data.Message) - - case webproto.TypeAOP: - var event aop.Event - if err := json.Unmarshal(message.Payload, &event); err != nil { - return output, events, fmt.Errorf("decode AOP payload: %w", err) - } - if !event.Valid() { - return output, events, fmt.Errorf("invalid AOP envelope") - } - events = append(events, event) - if monitor != nil { - monitor.renderEvent(event) - } - if sessionID == "" && event.Type == aop.TypeSessionStart { - sessionID = event.SessionID - } - if event.SessionID != sessionID { - continue - } - if turnID == "" && event.Type == aop.TypeTurnStart { - turnID = event.TurnID - } - if turnID != "" && event.TurnID != "" && event.TurnID != turnID { - continue - } - switch event.Type { - case aop.TypeMessage: - var data aop.MessageData - if json.Unmarshal(event.Data, &data) == nil && data.Role != "user" { - if text := messageText(data); text != "" { - output = text - } - } - case aop.TypeError: - var data aop.ErrorData - if json.Unmarshal(event.Data, &data) != nil || strings.TrimSpace(data.Message) == "" { - return output, events, fmt.Errorf("run AOP error has empty message") - } - agentErr = fmt.Errorf("agent error: %s", data.Message) - case aop.TypeTurnEnd: - var data aop.TurnEndData - if err := json.Unmarshal(event.Data, &data); err != nil { - return output, events, fmt.Errorf("decode turn.end: %w", err) - } - if agentErr != nil { - return output, events, agentErr - } - if data.Stop == "error" || data.Error != "" { - return output, events, fmt.Errorf("agent error: %s", strings.TrimSpace(data.Error)) - } - return output, events, nil - } - } - } -} diff --git a/core/harness/stdio_test.go b/core/harness/stdio_test.go deleted file mode 100644 index adb5139a..00000000 --- a/core/harness/stdio_test.go +++ /dev/null @@ -1,195 +0,0 @@ -//go:build e2e - -package harness - -import ( - "bytes" - "encoding/json" - "strings" - "testing" - - "github.com/chainreactors/aiscan/pkg/aop" - "github.com/chainreactors/aiscan/pkg/webproto" -) - -func TestConsumeAgentStream(t *testing.T) { - input := encodeFrames(t, - sessionOpenedFrame("root"), - aopFrame(aopTestEvent("root", "", aop.TypeSessionStart, aop.SessionStartData{})), - aopFrame(aopTestEvent("root", "turn-1", aop.TypeTurnStart, aop.TurnStartData{})), - aopFrame(aopMessageEvent("root", "turn-1", "assistant", "hello")), - aopFrame(aopTestEvent("root", "turn-1", aop.TypeTurnEnd, aop.TurnEndData{Stop: "completed"})), - ) - - var monitorOutput bytes.Buffer - output, events, err := consumeAgentStream(input, NewMonitor(&monitorOutput)) - if err != nil { - t.Fatalf("consumeAgentStream() error = %v", err) - } - if output != "hello" || len(events) != 4 { - t.Fatalf("output=%q events=%#v", output, events) - } - if !strings.Contains(monitorOutput.String(), "hello") || !strings.Contains(monitorOutput.String(), "run turn-1") { - t.Fatalf("monitor output = %q", monitorOutput.String()) - } -} - -func TestConsumeAgentStreamKeepsTypedToolData(t *testing.T) { - callID := "call-1" - input := encodeFrames(t, - sessionOpenedFrame("root"), - aopFrame(aopTestEvent("root", "turn-1", aop.TypeTurnStart, aop.TurnStartData{})), - aopFrame(aopTestEvent("root", "turn-1", aop.TypeToolCall, aop.ToolCallData{ - ToolCallID: callID, - ToolName: "bash", - Args: map[string]any{ - "command": "echo hello", - "nested": []any{map[string]any{"enabled": true}}, - }, - })), - aopFrame(aopTestEvent("root", "turn-1", aop.TypeToolResult, aop.ToolResultData{ - ToolCallID: callID, - ToolName: "bash", - Content: map[string]any{"output": []any{"hello", map[string]any{"code": float64(0)}}}, - })), - aopFrame(aopTestEvent("root", "turn-1", aop.TypeTurnEnd, aop.TurnEndData{Stop: "completed"})), - ) - - _, events, err := consumeAgentStream(input, nil) - if err != nil { - t.Fatalf("consumeAgentStream() error = %v", err) - } - calls := (&RunResult{Events: events}).ToolCalls() - if len(calls) != 1 { - t.Fatalf("tool calls = %#v", calls) - } - args, ok := calls[0].Args().(map[string]any) - if !ok || args["command"] != "echo hello" { - t.Fatalf("args = %#v", calls[0].Args()) - } - if !strings.Contains(calls[0].ResultText(), `"output":["hello"`) { - t.Fatalf("result = %q", calls[0].ResultText()) - } -} - -func TestConsumeAgentStreamWaitsForRootRunEnd(t *testing.T) { - input := encodeFrames(t, - sessionOpenedFrame("root"), - aopFrame(aopTestEvent("root", "turn-1", aop.TypeTurnStart, aop.TurnStartData{})), - aopFrame(aopTestEvent("child", "child-turn", aop.TypeTurnStart, aop.TurnStartData{})), - aopFrame(aopTestEvent("child", "child-turn", aop.TypeTurnEnd, aop.TurnEndData{Stop: "completed"})), - aopFrame(aopMessageEvent("root", "turn-1", "assistant", "root done")), - aopFrame(aopTestEvent("root", "turn-1", aop.TypeTurnEnd, aop.TurnEndData{Stop: "completed"})), - ) - output, events, err := consumeAgentStream(input, nil) - if err != nil || output != "root done" || len(events) != 5 { - t.Fatalf("output=%q events=%d err=%v", output, len(events), err) - } -} - -func TestConsumeAgentStreamReportsRootError(t *testing.T) { - input := encodeFrames(t, - sessionOpenedFrame("root"), - aopFrame(aopTestEvent("root", "turn-1", aop.TypeTurnStart, aop.TurnStartData{})), - aopFrame(aopTestEvent("root", "turn-1", aop.TypeError, aop.ErrorData{Message: "provider failed"})), - aopFrame(aopTestEvent("root", "turn-1", aop.TypeTurnEnd, aop.TurnEndData{Stop: "error", Error: "provider failed"})), - ) - _, events, err := consumeAgentStream(input, nil) - if err == nil || !strings.Contains(err.Error(), "provider failed") || len(events) != 3 { - t.Fatalf("events=%d error=%v", len(events), err) - } -} - -func TestConsumeAgentStreamReportsProtocolError(t *testing.T) { - input := encodeFrames(t, webproto.Message{ - Type: webproto.TypeError, - Payload: webproto.MustJSON(webproto.ErrorPayload{Message: "session rejected"}), - }) - _, _, err := consumeAgentStream(input, nil) - if err == nil || !strings.Contains(err.Error(), "session rejected") { - t.Fatalf("error = %v", err) - } -} - -func TestConsumeAgentStreamRejectsInvalidStreams(t *testing.T) { - tests := []struct { - name string - input *bytes.Buffer - needle string - }{ - { - name: "invalid envelope", - input: encodeFrames(t, - sessionOpenedFrame("root"), - webproto.Message{Type: webproto.TypeAOP, Payload: webproto.MustJSON(map[string]any{"type": "text"})}, - ), - needle: "invalid AOP envelope", - }, - { - name: "missing run terminal", - input: encodeFrames(t, - sessionOpenedFrame("root"), - aopFrame(aopTestEvent("root", "turn-1", aop.TypeTurnStart, aop.TurnStartData{})), - aopFrame(aopMessageEvent("root", "turn-1", "assistant", "hello")), - ), - needle: "without run turn.end", - }, - { - name: "no opened session", - input: encodeFrames(t), - needle: "without session.opened", - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - _, _, err := consumeAgentStream(tt.input, nil) - if err == nil || !strings.Contains(err.Error(), tt.needle) { - t.Fatalf("error = %v, want containing %q", err, tt.needle) - } - }) - } -} - -func aopTestEvent(sessionID, turnID, eventType string, data any) aop.Event { - raw, _ := json.Marshal(data) - return aop.Event{ - Type: eventType, - TS: "2026-07-19T00:00:00Z", - SessionID: sessionID, - TurnID: turnID, - Agent: "aiscan", - Data: raw, - } -} - -func aopMessageEvent(sessionID, turnID, role, text string) aop.Event { - return aopTestEvent(sessionID, turnID, aop.TypeMessage, aop.MessageData{ - MessageID: "m-1", - Role: role, - Parts: []aop.MessagePart{{Type: aop.PartText, Text: text}}, - }) -} - -func sessionOpenedFrame(sessionID string) webproto.Message { - return webproto.Message{ - Type: webproto.TypeSessionOpened, - Payload: webproto.MustJSON(webproto.SessionLifecyclePayload{SessionID: sessionID}), - } -} - -func aopFrame(event aop.Event) webproto.Message { - return webproto.Message{Type: webproto.TypeAOP, TurnID: event.TurnID, Payload: webproto.MustJSON(event)} -} - -func encodeFrames(t *testing.T, messages ...webproto.Message) *bytes.Buffer { - t.Helper() - var output bytes.Buffer - encoder := json.NewEncoder(&output) - for _, message := range messages { - if err := encoder.Encode(message); err != nil { - t.Fatal(err) - } - } - return &output -} diff --git a/core/harness/verify.go b/core/harness/verify.go deleted file mode 100644 index c903d22c..00000000 --- a/core/harness/verify.go +++ /dev/null @@ -1,278 +0,0 @@ -//go:build e2e - -package harness - -import ( - "fmt" - "strings" - "testing" - "time" -) - -// Verifier provides chainable assertions on a RunResult. -// Accumulates all failures; Done() reports them together. -// -// Two verification layers: -// -// Layer 1 — Structural (tool-level): -// -// Verify(t, r). -// OK(). -// Expect(Tool("bash").ArgContains("gogo").NoError()). -// Expect(Tool("subagent").Action("create").Arg("name", "worker")). -// Done() -// -// Layer 2 — Intent (outcome-level): -// -// Verify(t, r). -// OK(). -// ExpectInOrder( -// Tool("subagent").Action("create").Arg("name", "worker"), -// Tool("bash").ArgContains("scan"), -// ). -// OutputContains("worker"). -// NoToolErrors(). -// MaxTurns(5). -// Done() -type Verifier struct { - t *testing.T - r *RunResult - failures []string -} - -func Verify(t *testing.T, r *RunResult) *Verifier { - t.Helper() - return &Verifier{t: t, r: r} -} - -func (v *Verifier) fail(msg string) { v.failures = append(v.failures, msg) } - -func (v *Verifier) Done() { - v.t.Helper() - if len(v.failures) == 0 { - return - } - var sb strings.Builder - sb.WriteString(fmt.Sprintf("verification failed (%d issue(s)):\n", len(v.failures))) - for i, f := range v.failures { - sb.WriteString(fmt.Sprintf(" %d. %s\n", i+1, f)) - } - sb.WriteString(fmt.Sprintf("\nresult: exit=%d turns=%d tools=%d duration=%s\n", - v.r.ExitCode, v.r.Turns(), len(v.r.ToolCalls()), v.r.Duration)) - sb.WriteString(fmt.Sprintf("tool sequence: %v\n", v.r.ToolCallSequence())) - v.t.Fatal(sb.String()) -} - -// ===================================================================== -// Exit / Output -// ===================================================================== - -func (v *Verifier) OK() *Verifier { - if !v.r.OK() { - v.fail(fmt.Sprintf("exit code %d, expected 0\nstderr: %s", v.r.ExitCode, clip(v.r.Stderr, 500))) - } - return v -} - -func (v *Verifier) OutputContains(substr string) *Verifier { - if !v.r.ContainsOutput(substr) { - v.fail(fmt.Sprintf("output missing %q", substr)) - } - return v -} - -func (v *Verifier) OutputMissing(substr string) *Verifier { - if v.r.ContainsOutput(substr) { - v.fail(fmt.Sprintf("output should not contain %q", substr)) - } - return v -} - -// ===================================================================== -// Constraints -// ===================================================================== - -func (v *Verifier) MinTurns(n int) *Verifier { - if v.r.Turns() < n { - v.fail(fmt.Sprintf("expected >= %d turns, got %d", n, v.r.Turns())) - } - return v -} - -func (v *Verifier) MaxTurns(n int) *Verifier { - if v.r.Turns() > n { - v.fail(fmt.Sprintf("expected <= %d turns, got %d", n, v.r.Turns())) - } - return v -} - -func (v *Verifier) MinToolCalls(n int) *Verifier { - if len(v.r.ToolCalls()) < n { - v.fail(fmt.Sprintf("expected >= %d tool calls, got %d", n, len(v.r.ToolCalls()))) - } - return v -} - -func (v *Verifier) MaxToolCalls(n int) *Verifier { - if len(v.r.ToolCalls()) > n { - v.fail(fmt.Sprintf("expected <= %d tool calls, got %d", n, len(v.r.ToolCalls()))) - } - return v -} - -func (v *Verifier) CompletedWithin(d time.Duration) *Verifier { - if v.r.Duration > d { - v.fail(fmt.Sprintf("expected completion within %s, took %s", d, v.r.Duration)) - } - return v -} - -func (v *Verifier) ToolCount(name string, min, max int) *Verifier { - n := len(v.r.ToolCallsNamed(name)) - if n < min || n > max { - v.fail(fmt.Sprintf("tool %q called %d times, expected [%d, %d]", name, n, min, max)) - } - return v -} - -func (v *Verifier) ToolUsed(name string) *Verifier { - if !v.r.HasToolCall(name) { - v.fail(fmt.Sprintf("tool %q was not used", name)) - } - return v -} - -func (v *Verifier) ToolArgMatch(name string, match func(string) bool) *Verifier { - for _, call := range v.r.ToolCallsNamed(name) { - if match(argsText(call.Args())) { - return v - } - } - v.fail(fmt.Sprintf("no %q tool arguments matched", name)) - return v -} - -func (v *Verifier) ToolResultMatch(name string, match func(string) bool) *Verifier { - for _, call := range v.r.ToolCallsNamed(name) { - if match(call.ResultText()) { - return v - } - } - v.fail(fmt.Sprintf("no %q tool result matched", name)) - return v -} - -func (v *Verifier) AnyResultContains(substr string) *Verifier { - if !strings.Contains(v.r.AllToolResults(), substr) { - v.fail(fmt.Sprintf("no tool result contains %q", substr)) - } - return v -} - -// ===================================================================== -// Expect — pattern-based tool call verification -// ===================================================================== - -// Expect verifies that each pattern matches at least one tool call (any order). -func (v *Verifier) Expect(patterns ...ToolPattern) *Verifier { - result := matchUnordered(patterns, v.r.ToolCalls()) - for _, p := range result.unmatched { - v.fail(fmt.Sprintf("expected tool call not found: %s", p.describe())) - } - return v -} - -// ExpectInOrder verifies that patterns match tool calls in sequence -// (subsequence — other calls may appear between them). -func (v *Verifier) ExpectInOrder(patterns ...ToolPattern) *Verifier { - result := matchOrdered(patterns, v.r.ToolCalls()) - if len(result.unmatched) > 0 { - var descs []string - for _, p := range result.unmatched { - descs = append(descs, p.describe()) - } - v.fail(fmt.Sprintf("tool call sequence incomplete, unmatched: [%s]\nactual: %v", - strings.Join(descs, ", "), v.r.ToolCallSequence())) - } - return v -} - -// ExpectNone verifies that NO tool call matches the pattern. -func (v *Verifier) ExpectNone(patterns ...ToolPattern) *Verifier { - for _, p := range patterns { - for _, e := range v.r.ToolCalls() { - if p.Match(e) { - v.fail(fmt.Sprintf("unexpected tool call matched: %s", p.describe())) - break - } - } - } - return v -} - -// ===================================================================== -// Errors -// ===================================================================== - -func (v *Verifier) NoToolErrors() *Verifier { - errs := v.r.ErroredToolCalls() - if len(errs) > 0 { - names := make([]string, len(errs)) - for i, e := range errs { - names[i] = fmt.Sprintf("%s(%s)", e.Name(), clip(e.ResultText(), 80)) - } - v.fail(fmt.Sprintf("%d tool call(s) errored: %s", len(errs), strings.Join(names, ", "))) - } - return v -} - -// ===================================================================== -// Subagent shortcuts (built on Expect) -// ===================================================================== - -func (v *Verifier) SubagentCreated(name string) *Verifier { - return v.Expect(Tool("subagent").Arg("name", name)) -} - -func (v *Verifier) MinSubagentCreates(n int) *Verifier { - if v.r.SubagentCreateCount() < n { - v.fail(fmt.Sprintf("expected >= %d subagent creates, got %d", n, v.r.SubagentCreateCount())) - } - return v -} - -func (v *Verifier) SubagentResultContains(substr string) *Verifier { - for _, res := range v.r.SubagentResults() { - if strings.Contains(res, substr) { - return v - } - } - v.fail(fmt.Sprintf("no subagent result contains %q", substr)) - return v -} - -// ===================================================================== -// LLM Judge -// ===================================================================== - -// JudgeWith uses an LLM to evaluate whether the execution fulfilled the -// intent. The judge receives the full tool trace and final output, and -// returns a structured verdict. -// -// Verify(t, r). -// OK(). -// JudgeWith(h.Judge(), "create a loop, list it, delete it", ""). -// Done() -func (v *Verifier) JudgeWith(j *Judge, intent, criteria string) *Verifier { - verdict, err := j.Evaluate(intent, criteria, v.r) - if err != nil { - v.t.Logf("judge unavailable (degraded to warning): %s", err) - return v - } - v.t.Logf("judge: pass=%v score=%d reason=%q", verdict.Pass, verdict.Score, verdict.Reason) - if !verdict.Pass { - issues := strings.Join(verdict.Issues, "; ") - v.fail(fmt.Sprintf("judge failed (score=%d): %s [%s]", verdict.Score, verdict.Reason, issues)) - } - return v -} diff --git a/core/output/format.go b/core/output/format.go index 53de61cf..8cd3d6af 100644 --- a/core/output/format.go +++ b/core/output/format.go @@ -4,7 +4,7 @@ import ( "regexp" "strings" - "github.com/chainreactors/aiscan/pkg/agent/truncate" + "github.com/chainreactors/aiscan/core/truncate" ) var ansiPattern = regexp.MustCompile(`\x1b(?:\[[0-?]*[ -/]*[@-~]|\][^\x07]*(?:\x07|\x1b\\)|[PX^_].*?\x1b\\|[@-_])`) diff --git a/core/output/format_asset.go b/core/output/format_asset.go index 25f1c3a9..4d541f9f 100644 --- a/core/output/format_asset.go +++ b/core/output/format_asset.go @@ -1,288 +1,23 @@ package output import ( - "fmt" "net/url" - "sort" "strconv" "strings" ) +// FormatAssetReport renders the terminal asset report. The sitemap tree is the +// terminal report's own feature; everything else comes from the shared +// renderer in report.go. func FormatAssetReport(result *Result, color bool) string { - if result == nil { - return "Assets: 0 total\n" - } - c := NewColor(color) - - var sb strings.Builder - fmt.Fprintf(&sb, "Assets: %d total\n", len(result.Assets)) - fmt.Fprintf(&sb, "Summary: %d target(s), %d service(s), %d web endpoint(s), %d probe(s), %d loot(s), %d error(s), %s\n\n", - result.Summary.Targets, - result.Summary.Services, - result.Summary.Webs, - result.Summary.Probes, - result.Summary.Loots, - result.Summary.Errors, - result.Summary.Duration, - ) - - if len(result.Assets) == 0 { - return sb.String() - } - for i, asset := range result.Assets { - title := FirstNonEmpty(asset.Title, asset.Target, asset.Key) - fmt.Fprintf(&sb, "%d. %s\n", i+1, c.GreenBold(title)) - if asset.Target != "" && asset.Target != title { - fmt.Fprintf(&sb, " target: %s\n", asset.Target) - } - if asset.Status != "" { - fmt.Fprintf(&sb, " status: %s\n", asset.Status) - } - writeAssetTopItems(&sb, asset.Items, c) - writeAssetSitemap(&sb, asset, c) - if i < len(result.Assets)-1 { - sb.WriteByte('\n') - } - } - return sb.String() -} - -func writeAssetTopItems(sb *strings.Builder, items []AssetItem, c Color) { - for _, item := range items { - switch item.Kind { - case AssetItemPath: - continue - case AssetItemService: - line := strings.Join(CompactStrings( - AssetDataString(item.Data, "protocol"), - AssetDataString(item.Data, "service"), - AssetDataString(item.Data, "port"), - ), " ") - if line == "" { - line = FirstNonEmpty(item.Title, item.Target, item.Raw) - } - fmt.Fprintf(sb, " %s %s\n", c.Cyan("service:"), line) - case AssetItemFingerprint: - name := FirstNonEmpty(item.Title, item.Summary, item.Target) - fmt.Fprintf(sb, " %s %s\n", c.Cyan("fingerprint:"), name) - case AssetItemLoot, AssetItemNote, AssetItemResponse: - detail := AssetItemDetail(item) - line := FirstNonEmpty(item.Summary, item.Title, firstContentLine(detail), item.Raw) - if item.Status != "" { - line = c.Yellow("["+item.Status+"]") + " " + line - } - label := FirstNonEmpty(item.Source, item.Kind) - fmt.Fprintf(sb, " %s %s\n", c.Yellow(label+":"), line) - if detail != "" && detail != line && !strings.Contains(line, detail) { - for _, dl := range strings.Split(strings.TrimSpace(detail), "\n") { - if dl = strings.TrimSpace(dl); dl != "" { - fmt.Fprintf(sb, " %s\n", c.Dim(dl)) - } - } - } - case AssetItemError: - fmt.Fprintf(sb, " %s %s\n", c.Red("error:"), item.Summary) - } - } -} - -// --- sitemap rendering --- - -type sitemapEntry struct { - path string - status string - length int - title string - fingers []string - validated bool -} - -type sitemapNode struct { - segment string - status string - length int - title string - fingers []string - validated bool - isLeaf bool - annotations []string - children []*sitemapNode -} - -func writeAssetSitemap(sb *strings.Builder, asset Asset, c Color) { - var entries []sitemapEntry - for _, item := range asset.Items { - if item.Kind != AssetItemPath { - continue - } - p := FirstNonEmpty(AssetDataString(item.Data, "path"), WebPath(item.Target), item.Target) - if p == "" { - continue - } - entries = append(entries, sitemapEntry{ - path: p, - status: item.Status, - length: AssetDataInt(item.Data, "length"), - title: item.Title, - fingers: AssetDataStrings(item.Data, "fingers"), - validated: HasTag(item.Tags, "validated"), - }) - } - if len(entries) == 0 { - return - } - - sort.Slice(entries, func(i, j int) bool { return entries[i].path < entries[j].path }) - - sb.WriteString(" sitemap:\n") - root := buildSitemapTree(entries) - attachAnnotations(root, collectAnnotations(asset)) - renderNode(sb, root, " ", true, c) -} - -func buildSitemapTree(entries []sitemapEntry) *sitemapNode { - root := &sitemapNode{segment: "/"} - for _, e := range entries { - parts := splitPath(e.path) - if len(parts) == 0 { - root.isLeaf = true - root.status = e.status - root.length = e.length - root.title = e.title - root.fingers = mergeStrings(root.fingers, e.fingers) - root.validated = root.validated || e.validated - continue - } - node := root - for i, part := range parts { - child := findChild(node, part) - if child == nil { - child = &sitemapNode{segment: part} - node.children = append(node.children, child) - } - if i == len(parts)-1 { - child.isLeaf = true - child.status = e.status - child.length = e.length - child.title = e.title - child.fingers = mergeStrings(child.fingers, e.fingers) - child.validated = child.validated || e.validated - } - node = child - } - } - return root -} - -func collectAnnotations(asset Asset) map[string][]string { - out := make(map[string][]string) - for _, item := range asset.Items { - switch item.Kind { - case AssetItemFingerprint: - p := pathFromTarget(item.Target, asset.Target) - if p != "" { - out[p] = appendUniq(out[p], item.Title) - } - case AssetItemLoot, AssetItemNote, AssetItemResponse: - p := pathFromTarget(item.Target, asset.Target) - if p == "" { - p = "/" - } - skill := FirstNonEmpty(item.Source, item.Kind) - label := skill - if item.Status != "" { - label += ":" + item.Status - } - summary := FirstNonEmpty(item.Title, item.Summary) - if summary != "" && len(summary) <= 40 { - label += " " + summary - } - out[p] = appendUniq(out[p], label) - } - } - return out -} - -func attachAnnotations(root *sitemapNode, anns map[string][]string) { - if a, ok := anns["/"]; ok { - root.annotations = append(root.annotations, a...) - } - for path, a := range anns { - if path == "/" { - continue - } - parts := splitPath(path) - node := root - for _, part := range parts { - child := findChild(node, part) - if child == nil { - child = &sitemapNode{segment: part, isLeaf: true} - node.children = append(node.children, child) - } - node = child - } - node.annotations = append(node.annotations, a...) - } + return RenderReport(result, ReportOptions{ + Style: StyleANSI, + Color: color, + Sitemap: true, + }) } -func renderNode(sb *strings.Builder, node *sitemapNode, indent string, isRoot bool, c Color) { - var line strings.Builder - - if isRoot { - line.WriteString(indent) - } else { - line.WriteString(indent) - line.WriteString("├── ") - } - - if node.isLeaf && node.status != "" { - line.WriteString(c.Status(fmt.Sprintf("[%-3s]", node.status))) - } else { - line.WriteString(" ") - } - line.WriteString(" ") - - path := "/" + node.segment - if isRoot { - path = "/" - } - if node.validated { - line.WriteString(c.GreenBold(path)) - } else if node.isLeaf { - line.WriteString(path) - } else { - line.WriteString(c.Dim(path)) - } - - if node.isLeaf && node.length > 0 { - line.WriteString(" " + c.YellowBold(fmt.Sprintf("%d", node.length))) - } - - if node.title != "" && !isStaticTitle(node.title) { - line.WriteString(" " + c.Green(strconv.Quote(node.title))) - } - - if len(node.fingers) > 0 { - line.WriteString(" " + c.Cyan("["+strings.Join(node.fingers, ",")+"]")) - } - - for _, ann := range node.annotations { - line.WriteString(" " + c.Yellow("{"+ann+"}")) - } - - sb.WriteString(line.String()) - sb.WriteByte('\n') - - for _, child := range node.children { - childIndent := indent - if !isRoot { - childIndent += "│ " - } - renderNode(sb, child, childIndent, false, c) - } -} - -// --- shared helpers --- +// --- shared asset helpers --- func WebPath(rawURL string) string { parsed, err := url.Parse(strings.TrimSpace(rawURL)) @@ -381,74 +116,3 @@ func AssetDataStrings(data map[string]any, key string) []string { return nil } } - -func findChild(node *sitemapNode, segment string) *sitemapNode { - for _, c := range node.children { - if c.segment == segment { - return c - } - } - return nil -} - -func splitPath(p string) []string { - p = strings.Trim(p, "/") - if p == "" { - return nil - } - parts := strings.Split(p, "/") - if idx := strings.Index(parts[len(parts)-1], "?"); idx >= 0 { - parts[len(parts)-1] = parts[len(parts)-1][:idx] - } - return parts -} - -func pathFromTarget(target, assetTarget string) string { - if target == "" { - return "" - } - p := WebPath(target) - if p == target && assetTarget != "" { - if strings.HasPrefix(target, assetTarget) { - p = strings.TrimPrefix(target, assetTarget) - if p == "" { - p = "/" - } - } - } - return p -} - -func isStaticTitle(title string) bool { - switch strings.ToLower(title) { - case "js data", "css data", "ico data", "image data": - return true - } - return false -} - -func mergeStrings(a, b []string) []string { - if len(b) == 0 { - return a - } - seen := make(map[string]struct{}, len(a)) - for _, s := range a { - seen[strings.ToLower(s)] = struct{}{} - } - for _, s := range b { - if _, ok := seen[strings.ToLower(s)]; !ok { - a = append(a, s) - seen[strings.ToLower(s)] = struct{}{} - } - } - return a -} - -func appendUniq(slice []string, val string) []string { - for _, s := range slice { - if s == val { - return slice - } - } - return append(slice, val) -} diff --git a/core/output/report.go b/core/output/report.go new file mode 100644 index 00000000..12dd6d3b --- /dev/null +++ b/core/output/report.go @@ -0,0 +1,854 @@ +package output + +import ( + "fmt" + "sort" + "strconv" + "strings" + "time" +) + +// ReportStyle selects the emitter RenderReport hands the neutral model to. +type ReportStyle uint8 + +const ( + // StyleANSI is the operator-facing terminal report. + StyleANSI ReportStyle = iota + // StyleMarkdown is the report shipped to the web UI and to tool output. + StyleMarkdown +) + +// ReportOptions is the single knob set behind every asset report. The three +// renderers this replaced each owned a feature nobody else had (the sitemap +// tree, zh/en text + bare-host folding, the counter table), so the flags are +// what a caller opts into rather than what a style implies. +type ReportOptions struct { + Style ReportStyle + // Color enables ANSI escapes. StyleMarkdown ignores it — markdown output + // is never colorized. + Color bool + // Lang is "zh" or "en" (anything else, including "", means "en"). + // StyleANSI is English-only, so it ignores this. + Lang string + // Title is the report subject: the scan target for a web job, or a plain + // report name. Mode is the scan mode ("quick" / "full"); leaving Mode empty + // suppresses the target/mode/timestamp line and makes Title the bare H1. + // Markdown only. + Title string + Mode string + // Sitemap renders the per-asset path tree. + Sitemap bool + // CollapseBare folds live hosts that answered with nothing but non-web + // services into a trailing list instead of giving each one a section. + // Markdown only. + CollapseBare bool + // Metrics adds the counter table. Markdown only. + Metrics bool + // Inventory adds the flat per-kind sections (services / web evidence / + // findings / errors) the scan tool report carries. Markdown only. + Inventory bool +} + +// RenderReport walks Result → Asset → AssetItem exactly once into a neutral +// model, then emits it in the requested style. +func RenderReport(result *Result, opts ReportOptions) string { + model := buildReportModel(result, opts) + var report string + if opts.Style == StyleMarkdown { + report = renderMarkdownReport(model, opts) + } else { + report = renderANSIReport(model, opts) + } + return strings.TrimRight(report, " \t\r\n") + "\n" +} + +// --- neutral model --- + +type reportModel struct { + nilResult bool + summary Summary + total int + hosts int + fingers int + assets []reportAsset + bare []reportAsset +} + +type reportAsset struct { + title string // Title > Target > Key — the headline + label string // Target > Title > Key — the bare-host list entry + target string + status string + paths int + services []string + statuses []string + fingers []string + items []reportItem + sitemap *sitemapNode + isBare bool +} + +// reportItem is the per-item extraction — the part that used to exist in three +// places. text is the one-line rendering, name the short label used where a +// full line will not fit (sitemap annotations). +type reportItem struct { + kind string + label string // note-like items: Source > Kind + status string + target string + text string + name string + detail string + length int + fingers []string + validated bool +} + +func buildReportModel(result *Result, opts ReportOptions) reportModel { + if result == nil { + return reportModel{nilResult: true} + } + model := reportModel{summary: result.Summary, total: len(result.Assets)} + + hosts := make(map[string]struct{}) + fingers := make(map[string]struct{}) + for _, asset := range result.Assets { + item := buildReportAsset(asset, opts.Sitemap) + if host := reportAssetHost(asset); host != "" { + hosts[host] = struct{}{} + } + for _, finger := range item.fingers { + fingers[strings.ToLower(finger)] = struct{}{} + } + if opts.CollapseBare && item.isBare { + model.bare = append(model.bare, item) + continue + } + model.assets = append(model.assets, item) + } + + // An asset whose target parses to nothing still counts as a host. + model.hosts = len(hosts) + if model.hosts == 0 { + model.hosts = len(result.Assets) + } + model.fingers = len(fingers) + return model +} + +func buildReportAsset(asset Asset, sitemap bool) reportAsset { + out := reportAsset{ + title: FirstNonEmpty(asset.Title, asset.Target, asset.Key), + label: FirstNonEmpty(asset.Target, asset.Title, asset.Key), + target: asset.Target, + status: asset.Status, + } + + var services, statuses, fingers []string + annotations := make(map[string][]string) + hasService, onlyPlainServices := false, true + + for _, item := range asset.Items { + entry := reportItem{kind: item.Kind, status: item.Status, target: item.Target} + switch item.Kind { + case AssetItemService: + hasService = true + facts := strings.Join(CompactStrings( + AssetDataString(item.Data, "protocol"), + AssetDataString(item.Data, "service"), + AssetDataString(item.Data, "port"), + ), " ") + services = append(services, facts) + // A service with no structured facts still has a name to show. + entry.text = FirstNonEmpty(facts, item.Title, item.Target, item.Raw) + if isWebServiceItem(item) { + onlyPlainServices = false + } + case AssetItemFingerprint: + onlyPlainServices = false + entry.text = FirstNonEmpty(item.Title, item.Summary, AssetDataString(item.Data, "name"), item.Target) + entry.name = entry.text + fingers = append(fingers, entry.text) + if path := pathFromTarget(item.Target, asset.Target); path != "" { + annotations[path] = appendUniq(annotations[path], entry.text) + } + case AssetItemPath: + onlyPlainServices = false + out.paths++ + entry.text = FirstNonEmpty(AssetDataString(item.Data, "path"), WebPath(item.Target), item.Target) + entry.name = item.Title + entry.length = AssetDataInt(item.Data, "length") + entry.fingers = AssetDataStrings(item.Data, "fingers") + entry.validated = HasTag(item.Tags, "validated") + fingers = append(fingers, entry.fingers...) + if item.Status != "" { + statuses = append(statuses, item.Status) + } + case AssetItemLoot, AssetItemNote, AssetItemResponse, AssetItemError: + onlyPlainServices = false + entry.label = FirstNonEmpty(item.Source, item.Kind) + entry.detail = AssetItemDetail(item) + entry.text = FirstNonEmpty(item.Summary, item.Title, firstContentLine(entry.detail), item.Raw) + entry.name = FirstNonEmpty(item.Title, item.Summary) + if item.Kind != AssetItemError { + path := lootAnnotationPath(item, asset.Target) + annotations[path] = appendUniq(annotations[path], lootAnnotation(entry)) + } + default: + onlyPlainServices = false + entry.text = FirstNonEmpty(item.Summary, item.Title, item.Raw) + } + out.items = append(out.items, entry) + } + + out.isBare = hasService && onlyPlainServices + out.services = CompactStrings(services...) + out.statuses = CompactStrings(statuses...) + out.fingers = CompactStrings(fingers...) + if sitemap { + out.sitemap = buildSitemapTree(out.items, annotations) + } + return out +} + +// isWebServiceItem reports whether a service item is an HTTP-ish one, which is +// what keeps its host out of the "bare live host" bucket. +func isWebServiceItem(item AssetItem) bool { + svc := strings.ToLower(AssetDataString(item.Data, "service") + " " + AssetDataString(item.Data, "protocol")) + return strings.Contains(svc, "http") +} + +func lootAnnotationPath(item AssetItem, assetTarget string) string { + if path := pathFromTarget(item.Target, assetTarget); path != "" { + return path + } + return "/" +} + +// lootAnnotation is the compact "{skill:status summary}" tag hung off a +// sitemap node. Long summaries are dropped rather than wrapped. +func lootAnnotation(entry reportItem) string { + label := entry.label + if entry.status != "" { + label += ":" + entry.status + } + if entry.name != "" && len(entry.name) <= 40 { + label += " " + entry.name + } + return label +} + +// reportAssetHost reduces an asset to its host, so an IP that answered on both +// icmp and http counts once. +func reportAssetHost(asset Asset) string { + value := FirstNonEmpty(asset.Target, asset.Key, asset.Title) + if i := strings.Index(value, "://"); i >= 0 { + value = value[i+3:] + } + if i := strings.IndexAny(value, "/?#"); i >= 0 { + value = value[:i] + } + if strings.Count(value, ":") == 1 { // host:port — drop the port, leave IPv6 alone + value = value[:strings.LastIndex(value, ":")] + } + return value +} + +// --- ANSI emitter --- + +func renderANSIReport(model reportModel, opts ReportOptions) string { + if model.nilResult { + return "Assets: 0 total\n" + } + c := NewColor(opts.Color) + + var sb strings.Builder + fmt.Fprintf(&sb, "Assets: %d total\n", model.total) + fmt.Fprintf(&sb, "Summary: %d target(s), %d service(s), %d web endpoint(s), %d probe(s), %d loot(s), %d error(s), %s\n\n", + model.summary.Targets, + model.summary.Services, + model.summary.Webs, + model.summary.Probes, + model.summary.Loots, + model.summary.Errors, + model.summary.Duration, + ) + if model.total == 0 { + return sb.String() + } + + for i, asset := range model.assets { + fmt.Fprintf(&sb, "%d. %s\n", i+1, c.GreenBold(asset.title)) + if asset.target != "" && asset.target != asset.title { + fmt.Fprintf(&sb, " target: %s\n", asset.target) + } + if asset.status != "" { + fmt.Fprintf(&sb, " status: %s\n", asset.status) + } + for _, item := range asset.items { + writeANSIItem(&sb, item, c) + } + if asset.sitemap != nil { + sb.WriteString(" sitemap:\n") + renderSitemapNode(&sb, asset.sitemap, " ", true, c) + } + if i < len(model.assets)-1 { + sb.WriteByte('\n') + } + } + return sb.String() +} + +func writeANSIItem(sb *strings.Builder, item reportItem, c Color) { + switch item.kind { + case AssetItemPath: + return + case AssetItemService: + fmt.Fprintf(sb, " %s %s\n", c.Cyan("service:"), item.text) + case AssetItemFingerprint: + fmt.Fprintf(sb, " %s %s\n", c.Cyan("fingerprint:"), item.text) + case AssetItemLoot, AssetItemNote, AssetItemResponse: + line := item.text + if item.status != "" { + line = c.Yellow("["+item.status+"]") + " " + line + } + fmt.Fprintf(sb, " %s %s\n", c.Yellow(item.label+":"), line) + if item.detail != "" && item.detail != line && !strings.Contains(line, item.detail) { + for _, detailLine := range strings.Split(strings.TrimSpace(item.detail), "\n") { + if detailLine = strings.TrimSpace(detailLine); detailLine != "" { + fmt.Fprintf(sb, " %s\n", c.Dim(detailLine)) + } + } + } + case AssetItemError: + fmt.Fprintf(sb, " %s %s\n", c.Red("error:"), item.text) + } +} + +// --- markdown emitter --- + +// reportLang is the whole i18n surface: one flag, one lookup. +type reportLang struct{ zh bool } + +func newReportLang(lang string) reportLang { + return reportLang{zh: strings.HasPrefix(strings.ToLower(lang), "zh")} +} + +func (t reportLang) tr(zh, en string) string { + if t.zh { + return zh + } + return en +} + +func (t reportLang) sep() string { return t.tr("：", ": ") } + +func (t reportLang) modeName(mode string) string { + if strings.EqualFold(mode, "full") { + return t.tr("全面侦察", "Full recon") + } + return t.tr("快速侦察", "Quick recon") +} + +// renderMarkdownReport writes an operator-facing report: prose instead of a +// metric dump, no internal scanner names leaking into the text. +func renderMarkdownReport(model reportModel, opts ReportOptions) string { + t := newReportLang(opts.Lang) + + var sb strings.Builder + writeMarkdownHeader(&sb, t, opts) + if model.nilResult { + sb.WriteString(t.tr("本次扫描未返回结构化结果。\n", "No structured result was returned.\n")) + return sb.String() + } + + sb.WriteString("## " + t.tr("概述", "Overview") + "\n\n") + var overview strings.Builder + writeMarkdownOverview(&overview, t, model) + sb.WriteString(strings.TrimSpace(overview.String())) + sb.WriteString("\n\n") + + if opts.Metrics { + writeMarkdownMetrics(&sb, t, model) + } + if len(model.assets) > 0 { + sb.WriteString("## " + t.tr("资产明细", "Assets") + "\n\n") + for _, asset := range model.assets { + writeMarkdownAsset(&sb, t, asset, opts) + } + } + if len(model.bare) > 0 { + sb.WriteString("## " + t.tr("其他存活主机", "Other live hosts") + "\n\n") + for _, asset := range model.bare { + if len(asset.services) > 0 { + fmt.Fprintf(&sb, "- `%s` · %s\n", asset.label, strings.Join(asset.services, ", ")) + continue + } + fmt.Fprintf(&sb, "- `%s`\n", asset.label) + } + sb.WriteString("\n") + } + if opts.Inventory { + writeMarkdownInventory(&sb, t, model) + } + return sb.String() +} + +func writeMarkdownHeader(sb *strings.Builder, t reportLang, opts ReportOptions) { + if opts.Mode == "" { + fmt.Fprintf(sb, "# %s\n\n", FirstNonEmpty(opts.Title, t.tr("侦察报告", "Recon report"))) + sb.WriteString("---\n\n") + return + } + fmt.Fprintf(sb, "# %s%s\n\n", t.tr("侦察报告 · ", "Recon report · "), FirstNonEmpty(opts.Title, t.tr("目标", "target"))) + fmt.Fprintf(sb, "%s `%s` · %s · %s\n\n", + t.tr("目标", "Target"), opts.Title, + t.modeName(opts.Mode), + time.Now().Format("2006-01-02 15:04:05")) + sb.WriteString("---\n\n") +} + +// writeMarkdownOverview is the executive summary — one flowing paragraph that +// names only the numbers actually present, so a clean scan reads like a +// sentence rather than a table full of zeros. +func writeMarkdownOverview(sb *strings.Builder, t reportLang, model reportModel) { + s := model.summary + if t.zh { + fmt.Fprintf(sb, "本次侦察共识别 %d 台主机、%d 个开放服务", model.hosts, s.Services) + if s.Webs > 0 { + fmt.Fprintf(sb, "（含 %d 个 Web 站点）", s.Webs) + } + sb.WriteString("。") + if s.Probes > 0 { + fmt.Fprintf(sb, "累计探测 %d 条路径", s.Probes) + if model.fingers > 0 { + fmt.Fprintf(sb, "、命中 %d 项 Web 指纹", model.fingers) + } + sb.WriteString("。") + } else if model.fingers > 0 { + fmt.Fprintf(sb, "命中 %d 项 Web 指纹。", model.fingers) + } + if s.Loots > 0 { + fmt.Fprintf(sb, "**发现 %d 项需优先复核的安全发现（凭证 / 弱口令 / 漏洞）。**", s.Loots) + } + if s.Errors > 0 { + fmt.Fprintf(sb, "另有 %d 处探测报错。", s.Errors) + } + if s.Duration != "" { + fmt.Fprintf(sb, "全程耗时 %s。", s.Duration) + } + return + } + + fmt.Fprintf(sb, "The scan identified %s across %s", reportPlural(model.hosts, "host", "hosts"), reportPlural(s.Services, "open service", "open services")) + if s.Webs > 0 { + fmt.Fprintf(sb, " (%s)", reportPlural(s.Webs, "web site", "web sites")) + } + sb.WriteString(". ") + if s.Probes > 0 { + fmt.Fprintf(sb, "It probed %s", reportPlural(s.Probes, "path", "paths")) + if model.fingers > 0 { + fmt.Fprintf(sb, " and matched %s", reportPlural(model.fingers, "fingerprint", "fingerprints")) + } + sb.WriteString(". ") + } else if model.fingers > 0 { + fmt.Fprintf(sb, "It matched %s. ", reportPlural(model.fingers, "fingerprint", "fingerprints")) + } + if s.Loots > 0 { + fmt.Fprintf(sb, "**%s surfaced (credentials / weak passwords / vulnerabilities) — review these first.** ", reportPlural(s.Loots, "security finding", "security findings")) + } + if s.Errors > 0 { + fmt.Fprintf(sb, "%s occurred during probing. ", reportPlural(s.Errors, "error", "errors")) + } + if s.Duration != "" { + fmt.Fprintf(sb, "The scan took %s.", s.Duration) + } +} + +func writeMarkdownMetrics(sb *strings.Builder, t reportLang, model reportModel) { + s := model.summary + sb.WriteString("## " + t.tr("指标", "Metrics") + "\n\n") + fmt.Fprintf(sb, "| %s | %s |\n", t.tr("指标", "Metric"), t.tr("数值", "Value")) + sb.WriteString("| --- | ---: |\n") + for _, row := range []struct { + label string + value any + }{ + {t.tr("输入目标", "Inputs"), s.Targets}, + {t.tr("开放服务", "Open services"), s.Services}, + {t.tr("Web 站点", "Web endpoints"), s.Webs}, + {t.tr("路径探测", "Web probes"), s.Probes}, + {t.tr("Web 指纹", "Fingerprints"), model.fingers}, + {t.tr("安全发现", "Loots"), s.Loots}, + {t.tr("错误", "Errors"), s.Errors}, + {t.tr("任务", "Tasks"), s.Tasks}, + {t.tr("请求", "Requests"), s.Requests}, + {t.tr("耗时", "Duration"), s.Duration}, + } { + fmt.Fprintf(sb, "| %s | %v |\n", row.label, row.value) + } + sb.WriteString("\n") +} + +func writeMarkdownAsset(sb *strings.Builder, t reportLang, asset reportAsset, opts ReportOptions) { + title := FirstNonEmpty(asset.title, t.tr("资产", "Asset")) + if asset.target != "" && asset.target != title { + fmt.Fprintf(sb, "### %s — `%s`\n\n", title, asset.target) + } else { + fmt.Fprintf(sb, "### %s\n\n", title) + } + + writeMarkdownFact(sb, t, t.tr("开放服务", "Services"), asset.services) + writeMarkdownFact(sb, t, t.tr("HTTP 响应", "HTTP"), asset.statuses) + writeMarkdownFact(sb, t, t.tr("Web 指纹", "Fingerprints"), asset.fingers) + if asset.paths > 0 { + fmt.Fprintf(sb, "- %s%s%s\n", t.tr("已探测路径", "Paths"), t.sep(), t.tr(fmt.Sprintf("%d 条", asset.paths), strconv.Itoa(asset.paths))) + } + if asset.status != "" { + fmt.Fprintf(sb, "- %s%s%s\n", t.tr("状态", "State"), t.sep(), markdownCode(asset.status)) + } + sb.WriteString("\n") + + if opts.Sitemap && asset.sitemap != nil { + sb.WriteString("#### " + t.tr("站点地图", "Sitemap") + "\n\n```text\n") + renderSitemapNode(sb, asset.sitemap, "", true, NewColor(false)) + sb.WriteString("```\n\n") + } + writeMarkdownAnalysis(sb, t, asset.items) +} + +func writeMarkdownFact(sb *strings.Builder, t reportLang, label string, values []string) { + if len(values) == 0 { + return + } + coded := make([]string, 0, len(values)) + for _, value := range values { + coded = append(coded, markdownCode(value)) + } + fmt.Fprintf(sb, "- %s%s%s\n", label, t.sep(), strings.Join(coded, t.tr("、", ", "))) +} + +func writeMarkdownAnalysis(sb *strings.Builder, t reportLang, items []reportItem) { + wrote := false + for _, item := range items { + switch item.kind { + case AssetItemLoot, AssetItemNote, AssetItemResponse, AssetItemError: + default: + continue + } + if item.text == "" { + continue + } + if !wrote { + sb.WriteString("#### " + t.tr("分析研判", "Analysis") + "\n\n") + wrote = true + } + fmt.Fprintf(sb, "##### %s\n\n", markdownHeading(item.text)) + switch { + case item.detail != "" && strings.TrimSpace(item.text) != strings.TrimSpace(item.detail): + sb.WriteString(item.detail) + sb.WriteString("\n\n") + case item.detail == "": + sb.WriteString(item.text) + sb.WriteString("\n\n") + } + } +} + +// writeMarkdownInventory is the flat cross-asset listing the scan tool report +// has always carried: every service, probe, finding and error in one place. +func writeMarkdownInventory(sb *strings.Builder, t reportLang, model reportModel) { + assets := make([]reportAsset, 0, len(model.assets)+len(model.bare)) + assets = append(assets, model.assets...) + assets = append(assets, model.bare...) + + var services, paths, findings, errors []string + for _, asset := range assets { + for _, item := range asset.items { + switch item.kind { + case AssetItemService: + services = append(services, fmt.Sprintf("- %s · %s\n", + markdownCode(FirstNonEmpty(item.target, asset.label)), item.text)) + case AssetItemPath: + paths = append(paths, "- "+strings.Join(pathInventoryParts(item), " · ")+"\n") + case AssetItemLoot, AssetItemNote, AssetItemResponse: + findings = append(findings, markdownStatusLine(findingInventoryLine(item), item.status)) + case AssetItemError: + errors = append(errors, "- "+item.text+"\n") + } + } + } + + writeMarkdownSection(sb, t.tr("开放服务", "Open Services"), services) + writeMarkdownSection(sb, t.tr("Web 证据", "Web Evidence"), paths) + writeMarkdownSection(sb, t.tr("安全发现", "Findings"), findings) + writeMarkdownSection(sb, t.tr("错误", "Errors"), errors) +} + +func pathInventoryParts(item reportItem) []string { + parts := []string{markdownCode(FirstNonEmpty(item.target, item.text))} + if item.status != "" { + parts = append(parts, markdownCode(item.status)) + } + if item.name != "" && !isStaticTitle(item.name) { + parts = append(parts, strconv.Quote(item.name)) + } + if len(item.fingers) > 0 { + parts = append(parts, markdownCode(strings.Join(item.fingers, ","))) + } + return parts +} + +func findingInventoryLine(item reportItem) string { + line := item.text + if item.target != "" { + line += " — " + markdownCode(item.target) + } + return line +} + +// markdownStatusLine carries the verification verdict into the bullet, so an +// unconfirmed finding cannot be mistaken for a proven one. +func markdownStatusLine(line, status string) string { + if line == "" { + return "" + } + switch status { + case "not_confirmed": + return "- ~~" + line + "~~ *(not confirmed)*\n" + case "confirmed": + return "- **[verified]** " + line + "\n" + case "inconclusive": + return "- **[inconclusive]** " + line + "\n" + case "failed": + return "- **[verification failed]** " + line + "\n" + default: + return "- " + line + "\n" + } +} + +func writeMarkdownSection(sb *strings.Builder, heading string, lines []string) { + if len(lines) == 0 { + return + } + sb.WriteString("## " + heading + "\n\n") + for _, line := range lines { + sb.WriteString(line) + } + sb.WriteString("\n") +} + +func markdownCode(value string) string { + value = strings.ReplaceAll(value, "`", "'") + return "`" + value + "`" +} + +func markdownHeading(value string) string { + value = strings.TrimSpace(value) + value = strings.ReplaceAll(value, "\n", " ") + if value == "" { + return "Analysis" + } + return strings.TrimLeft(value, "# ") +} + +func reportPlural(n int, one, many string) string { + if n == 1 { + return fmt.Sprintf("%d %s", n, one) + } + return fmt.Sprintf("%d %s", n, many) +} + +// --- sitemap tree --- + +type sitemapNode struct { + segment string + status string + length int + title string + fingers []string + validated bool + isLeaf bool + annotations []string + children []*sitemapNode +} + +// buildSitemapTree folds the asset's path items into a directory tree and hangs +// the fingerprint / finding annotations off the node they were found on. +// Returns nil when the asset has no paths, so callers can skip the section. +func buildSitemapTree(items []reportItem, annotations map[string][]string) *sitemapNode { + paths := make([]reportItem, 0, len(items)) + for _, item := range items { + if item.kind == AssetItemPath && item.text != "" { + paths = append(paths, item) + } + } + if len(paths) == 0 { + return nil + } + sort.Slice(paths, func(i, j int) bool { return paths[i].text < paths[j].text }) + + root := &sitemapNode{segment: "/"} + for _, item := range paths { + node := root + for _, part := range splitPath(item.text) { + child := findSitemapChild(node, part) + if child == nil { + child = &sitemapNode{segment: part} + node.children = append(node.children, child) + } + node = child + } + node.isLeaf = true + node.status = item.status + node.length = item.length + node.title = item.name + node.fingers = mergeStrings(node.fingers, item.fingers) + node.validated = node.validated || item.validated + } + attachSitemapAnnotations(root, annotations) + return root +} + +func attachSitemapAnnotations(root *sitemapNode, annotations map[string][]string) { + if values, ok := annotations["/"]; ok { + root.annotations = append(root.annotations, values...) + } + for path, values := range annotations { + if path == "/" { + continue + } + node := root + for _, part := range splitPath(path) { + child := findSitemapChild(node, part) + if child == nil { + child = &sitemapNode{segment: part, isLeaf: true} + node.children = append(node.children, child) + } + node = child + } + node.annotations = append(node.annotations, values...) + } +} + +func renderSitemapNode(sb *strings.Builder, node *sitemapNode, indent string, isRoot bool, c Color) { + var line strings.Builder + + line.WriteString(indent) + if !isRoot { + line.WriteString("├── ") + } + + if node.isLeaf && node.status != "" { + line.WriteString(c.Status(fmt.Sprintf("[%-3s]", node.status))) + } else { + line.WriteString(" ") + } + line.WriteString(" ") + + path := "/" + node.segment + if isRoot { + path = "/" + } + switch { + case node.validated: + line.WriteString(c.GreenBold(path)) + case node.isLeaf: + line.WriteString(path) + default: + line.WriteString(c.Dim(path)) + } + + if node.isLeaf && node.length > 0 { + line.WriteString(" " + c.YellowBold(strconv.Itoa(node.length))) + } + if node.title != "" && !isStaticTitle(node.title) { + line.WriteString(" " + c.Green(strconv.Quote(node.title))) + } + if len(node.fingers) > 0 { + line.WriteString(" " + c.Cyan("["+strings.Join(node.fingers, ",")+"]")) + } + for _, annotation := range node.annotations { + line.WriteString(" " + c.Yellow("{"+annotation+"}")) + } + + sb.WriteString(line.String()) + sb.WriteByte('\n') + + for _, child := range node.children { + childIndent := indent + if !isRoot { + childIndent += "│ " + } + renderSitemapNode(sb, child, childIndent, false, c) + } +} + +func findSitemapChild(node *sitemapNode, segment string) *sitemapNode { + for _, child := range node.children { + if child.segment == segment { + return child + } + } + return nil +} + +func splitPath(p string) []string { + p = strings.Trim(p, "/") + if p == "" { + return nil + } + parts := strings.Split(p, "/") + if idx := strings.Index(parts[len(parts)-1], "?"); idx >= 0 { + parts[len(parts)-1] = parts[len(parts)-1][:idx] + } + return parts +} + +func pathFromTarget(target, assetTarget string) string { + if target == "" { + return "" + } + p := WebPath(target) + if p == target && assetTarget != "" && strings.HasPrefix(target, assetTarget) { + p = strings.TrimPrefix(target, assetTarget) + if p == "" { + p = "/" + } + } + return p +} + +func isStaticTitle(title string) bool { + switch strings.ToLower(title) { + case "js data", "css data", "ico data", "image data": + return true + } + return false +} + +func mergeStrings(a, b []string) []string { + if len(b) == 0 { + return a + } + seen := make(map[string]struct{}, len(a)) + for _, s := range a { + seen[strings.ToLower(s)] = struct{}{} + } + for _, s := range b { + if _, ok := seen[strings.ToLower(s)]; !ok { + a = append(a, s) + seen[strings.ToLower(s)] = struct{}{} + } + } + return a +} + +func appendUniq(slice []string, val string) []string { + for _, s := range slice { + if s == val { + return slice + } + } + return append(slice, val) +} diff --git a/core/output/report_golden_test.go b/core/output/report_golden_test.go new file mode 100644 index 00000000..014cc660 --- /dev/null +++ b/core/output/report_golden_test.go @@ -0,0 +1,113 @@ +package output + +import ( + "encoding/json" + "flag" + "os" + "path/filepath" + "regexp" + "testing" +) + +// updateReportGolden rewrites the .golden files instead of comparing against +// them, so the diff of a deliberate rendering change is reviewable on its own. +var updateReportGolden = flag.Bool("update-report-golden", false, "rewrite report golden files") + +// LoadReportFixture reads one of the shared report fixtures. It lives in +// core/output because that is where the fixtures live, but pkg/web reads the +// same files so the two renderers are pinned against identical input. +func loadReportFixture(t *testing.T, name string) *Result { + t.Helper() + raw, err := os.ReadFile(filepath.Join("testdata", name+".json")) + if err != nil { + t.Fatalf("read fixture: %v", err) + } + result := &Result{} + if err := json.Unmarshal(raw, result); err != nil { + t.Fatalf("decode fixture: %v", err) + } + return result +} + +// reportStamp matches the "generated at" header timestamp, the one part of the +// markdown report that cannot be pinned. +var reportStamp = regexp.MustCompile(`\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}`) + +func checkReportGolden(t *testing.T, name, got string) { + t.Helper() + got = reportStamp.ReplaceAllString(got, "") + path := filepath.Join("testdata", name+".golden") + if *updateReportGolden { + if err := os.WriteFile(path, []byte(got), 0o644); err != nil { + t.Fatalf("write golden: %v", err) + } + return + } + want, err := os.ReadFile(path) + if err != nil { + t.Fatalf("read golden (run go test -run %s -update-report-golden): %v", t.Name(), err) + } + if got != string(want) { + t.Errorf("%s mismatch\n--- got ---\n%s\n--- want ---\n%s", path, got, string(want)) + } +} + +// TestAssetReportGolden pins the terminal asset report. FormatAssetReport is a +// wrapper over RenderReport, so this is also the ANSI emitter's contract. +func TestAssetReportGolden(t *testing.T) { + for _, tc := range []struct { + name string + fixture string + color bool + }{ + {name: "asset_plain", fixture: "report_fixture"}, + {name: "asset_color", fixture: "report_fixture", color: true}, + {name: "asset_empty", fixture: "report_empty"}, + } { + t.Run(tc.name, func(t *testing.T) { + checkReportGolden(t, tc.name, FormatAssetReport(loadReportFixture(t, tc.fixture), tc.color)) + }) + } +} + +// TestRenderReportMarkdownGolden pins the markdown emitter under the option +// sets its two callers use. The web_* goldens must stay byte-identical to +// pkg/web/testdata/web_*.golden — that is the cross-package parity check. +func TestRenderReportMarkdownGolden(t *testing.T) { + web := ReportOptions{Style: StyleMarkdown, Title: "10.0.0.1", Mode: "quick", CollapseBare: true} + tool := ReportOptions{Style: StyleMarkdown, Title: "Scan Report", Sitemap: true, CollapseBare: true, Metrics: true, Inventory: true} + + for _, tc := range []struct { + name string + fixture string + opts ReportOptions + nilRes bool + }{ + {name: "md_web_zh", fixture: "report_fixture", opts: withLang(web, "zh")}, + {name: "md_web_en", fixture: "report_fixture", opts: withLang(web, "en")}, + {name: "md_web_empty_zh", fixture: "report_empty", opts: withLang(web, "zh")}, + {name: "md_web_empty_en", fixture: "report_empty", opts: withLang(web, "en")}, + {name: "md_web_nil", opts: withLang(web, "en"), nilRes: true}, + {name: "md_tool", fixture: "report_fixture", opts: tool}, + {name: "md_tool_empty", fixture: "report_empty", opts: tool}, + } { + t.Run(tc.name, func(t *testing.T) { + var result *Result + if !tc.nilRes { + result = loadReportFixture(t, tc.fixture) + } + checkReportGolden(t, tc.name, RenderReport(result, tc.opts)) + }) + } +} + +func withLang(opts ReportOptions, lang string) ReportOptions { + opts.Lang = lang + return opts +} + +func TestAssetReportNilResult(t *testing.T) { + if got := FormatAssetReport(nil, false); got != "Assets: 0 total\n" { + t.Fatalf("nil result = %q", got) + } +} diff --git a/core/output/testdata/asset_color.golden b/core/output/testdata/asset_color.golden new file mode 100644 index 00000000..144a4085 --- /dev/null +++ b/core/output/testdata/asset_color.golden @@ -0,0 +1,43 @@ +Assets: 4 total +Summary: 2 target(s), 4 service(s), 2 web endpoint(s), 3 probe(s), 3 loot(s), 1 error(s), 22.266s + +1. Example App + target: http://10.0.0.1 + status: confirmed + service: tcp http 80 + fingerprint: nginx + vuln: [confirmed] CVE-2021-41773 path traversal + ## Impact + Arbitrary file read through `/cgi-bin/`. + | Field | Value | + |---|---| + | CVSS | 9.8 | + weakpass: [high] ssh root:root + deep: [info] Admin console is reachable without authentication. + deep: [response] Let me analyze the collected browser evidence. + Let me analyze the collected browser evidence. + ## Evidence Analysis + | Asset | Details | + |---|---| + | API | GET /api/scans | + error: dial tcp 10.0.0.1:8443: connect: connection refused + sitemap: + [200] / 1256 "Example App" [nginx] {nginx} {deep:response} + ├── /admin {vuln:confirmed CVE-2021-41773 path traversal} {deep:info} + │ ├── [401] /login 512 "Login" [basic-auth] + ├── /static + │ ├── [200] /app.js 9001 + ├── /10.0.0.1:22 {weakpass:high ssh root:root} + +2. MySQL 5.7.32 + target: 10.0.0.2:3306 + status: loot + service: tcp mysql 3306 + fingerprint: [loot] mysql 5.7.32 + +3. icmp + target: 10.0.0.2:icmp + service: icmp + +4. 10.0.0.3:445 + service: tcp smb 445 diff --git a/core/output/testdata/asset_empty.golden b/core/output/testdata/asset_empty.golden new file mode 100644 index 00000000..90553463 --- /dev/null +++ b/core/output/testdata/asset_empty.golden @@ -0,0 +1,2 @@ +Assets: 0 total +Summary: 0 target(s), 0 service(s), 0 web endpoint(s), 0 probe(s), 0 loot(s), 0 error(s), diff --git a/core/output/testdata/asset_plain.golden b/core/output/testdata/asset_plain.golden new file mode 100644 index 00000000..703a0faf --- /dev/null +++ b/core/output/testdata/asset_plain.golden @@ -0,0 +1,43 @@ +Assets: 4 total +Summary: 2 target(s), 4 service(s), 2 web endpoint(s), 3 probe(s), 3 loot(s), 1 error(s), 22.266s + +1. Example App + target: http://10.0.0.1 + status: confirmed + service: tcp http 80 + fingerprint: nginx + vuln: [confirmed] CVE-2021-41773 path traversal + ## Impact + Arbitrary file read through `/cgi-bin/`. + | Field | Value | + |---|---| + | CVSS | 9.8 | + weakpass: [high] ssh root:root + deep: [info] Admin console is reachable without authentication. + deep: [response] Let me analyze the collected browser evidence. + Let me analyze the collected browser evidence. + ## Evidence Analysis + | Asset | Details | + |---|---| + | API | GET /api/scans | + error: dial tcp 10.0.0.1:8443: connect: connection refused + sitemap: + [200] / 1256 "Example App" [nginx] {nginx} {deep:response} + ├── /admin {vuln:confirmed CVE-2021-41773 path traversal} {deep:info} + │ ├── [401] /login 512 "Login" [basic-auth] + ├── /static + │ ├── [200] /app.js 9001 + ├── /10.0.0.1:22 {weakpass:high ssh root:root} + +2. MySQL 5.7.32 + target: 10.0.0.2:3306 + status: loot + service: tcp mysql 3306 + fingerprint: [loot] mysql 5.7.32 + +3. icmp + target: 10.0.0.2:icmp + service: icmp + +4. 10.0.0.3:445 + service: tcp smb 445 diff --git a/core/output/testdata/md_tool.golden b/core/output/testdata/md_tool.golden new file mode 100644 index 00000000..8b197cd1 --- /dev/null +++ b/core/output/testdata/md_tool.golden @@ -0,0 +1,116 @@ +# Scan Report + +--- + +## Overview + +The scan identified 3 hosts across 4 open services (2 web sites). It probed 3 paths and matched 2 fingerprints. **3 security findings surfaced (credentials / weak passwords / vulnerabilities) — review these first.** 1 error occurred during probing. The scan took 22.266s. + +## Metrics + +| Metric | Value | +| --- | ---: | +| Inputs | 2 | +| Open services | 4 | +| Web endpoints | 2 | +| Web probes | 3 | +| Fingerprints | 2 | +| Loots | 3 | +| Errors | 1 | +| Tasks | 7 | +| Requests | 19 | +| Duration | 22.266s | + +## Assets + +### Example App — `http://10.0.0.1` + +- Services: `tcp http 80` +- HTTP: `200`, `401` +- Fingerprints: `nginx`, `basic-auth` +- Paths: 3 +- State: `confirmed` + +#### Sitemap + +```text +[200] / 1256 "Example App" [nginx] {nginx} {deep:response} +├── /admin {vuln:confirmed CVE-2021-41773 path traversal} {deep:info} +│ ├── [401] /login 512 "Login" [basic-auth] +├── /static +│ ├── [200] /app.js 9001 +├── /10.0.0.1:22 {weakpass:high ssh root:root} +``` + +#### Analysis + +##### CVE-2021-41773 path traversal + +## Impact + +Arbitrary file read through `/cgi-bin/`. + +| Field | Value | +|---|---| +| CVSS | 9.8 | + +##### ssh root:root + +ssh root:root + +##### Admin console is reachable without authentication. + +##### Let me analyze the collected browser evidence. + +Let me analyze the collected browser evidence. + +## Evidence Analysis + +| Asset | Details | +|---|---| +| API | GET /api/scans | + +##### dial tcp 10.0.0.1:8443: connect: connection refused + +dial tcp 10.0.0.1:8443: connect: connection refused + +### MySQL 5.7.32 — `10.0.0.2:3306` + +- Services: `tcp mysql 3306` +- State: `loot` + +#### Analysis + +##### mysql 5.7.32 + +mysql 5.7.32 + +## Other live hosts + +- `10.0.0.2:icmp` · icmp +- `10.0.0.3:445` · tcp smb 445 + +## Open Services + +- `10.0.0.1:80` · tcp http 80 +- `10.0.0.2:3306` · tcp mysql 3306 +- `10.0.0.2:icmp` · icmp +- `10.0.0.3:445` · tcp smb 445 + +## Web Evidence + +- `http://10.0.0.1/` · `200` · "Example App" · `nginx` +- `http://10.0.0.1/admin/login` · `401` · "Login" · `basic-auth` +- `http://10.0.0.1/static/app.js` · `200` + +## Findings + +- **[verified]** CVE-2021-41773 path traversal — `http://10.0.0.1/admin` +- ssh root:root — `10.0.0.1:22` +- Admin console is reachable without authentication. — `http://10.0.0.1/admin` +- Let me analyze the collected browser evidence. — `http://10.0.0.1` +- mysql 5.7.32 — `10.0.0.2:3306` + +## Errors + +- dial tcp 10.0.0.1:8443: connect: connection refused diff --git a/core/output/testdata/md_tool_empty.golden b/core/output/testdata/md_tool_empty.golden new file mode 100644 index 00000000..a016ccdc --- /dev/null +++ b/core/output/testdata/md_tool_empty.golden @@ -0,0 +1,22 @@ +# Scan Report + +--- + +## Overview + +The scan identified 0 hosts across 0 open services. + +## Metrics + +| Metric | Value | +| --- | ---: | +| Inputs | 0 | +| Open services | 0 | +| Web endpoints | 0 | +| Web probes | 0 | +| Fingerprints | 0 | +| Loots | 0 | +| Errors | 0 | +| Tasks | 0 | +| Requests | 0 | +| Duration | | diff --git a/core/output/testdata/md_web_empty_en.golden b/core/output/testdata/md_web_empty_en.golden new file mode 100644 index 00000000..99152d2b --- /dev/null +++ b/core/output/testdata/md_web_empty_en.golden @@ -0,0 +1,9 @@ +# Recon report · 10.0.0.1 + +Target `10.0.0.1` · Quick recon · + +--- + +## Overview + +The scan identified 0 hosts across 0 open services. diff --git a/core/output/testdata/md_web_empty_zh.golden b/core/output/testdata/md_web_empty_zh.golden new file mode 100644 index 00000000..ef7f420d --- /dev/null +++ b/core/output/testdata/md_web_empty_zh.golden @@ -0,0 +1,9 @@ +# 侦察报告 · 10.0.0.1 + +目标 `10.0.0.1` · 快速侦察 · + +--- + +## 概述 + +本次侦察共识别 0 台主机、0 个开放服务。 diff --git a/core/output/testdata/md_web_en.golden b/core/output/testdata/md_web_en.golden new file mode 100644 index 00000000..178138df --- /dev/null +++ b/core/output/testdata/md_web_en.golden @@ -0,0 +1,67 @@ +# Recon report · 10.0.0.1 + +Target `10.0.0.1` · Quick recon · + +--- + +## Overview + +The scan identified 3 hosts across 4 open services (2 web sites). It probed 3 paths and matched 2 fingerprints. **3 security findings surfaced (credentials / weak passwords / vulnerabilities) — review these first.** 1 error occurred during probing. The scan took 22.266s. + +## Assets + +### Example App — `http://10.0.0.1` + +- Services: `tcp http 80` +- HTTP: `200`, `401` +- Fingerprints: `nginx`, `basic-auth` +- Paths: 3 +- State: `confirmed` + +#### Analysis + +##### CVE-2021-41773 path traversal + +## Impact + +Arbitrary file read through `/cgi-bin/`. + +| Field | Value | +|---|---| +| CVSS | 9.8 | + +##### ssh root:root + +ssh root:root + +##### Admin console is reachable without authentication. + +##### Let me analyze the collected browser evidence. + +Let me analyze the collected browser evidence. + +## Evidence Analysis + +| Asset | Details | +|---|---| +| API | GET /api/scans | + +##### dial tcp 10.0.0.1:8443: connect: connection refused + +dial tcp 10.0.0.1:8443: connect: connection refused + +### MySQL 5.7.32 — `10.0.0.2:3306` + +- Services: `tcp mysql 3306` +- State: `loot` + +#### Analysis + +##### mysql 5.7.32 + +mysql 5.7.32 + +## Other live hosts + +- `10.0.0.2:icmp` · icmp +- `10.0.0.3:445` · tcp smb 445 diff --git a/core/output/testdata/md_web_nil.golden b/core/output/testdata/md_web_nil.golden new file mode 100644 index 00000000..15413fd8 --- /dev/null +++ b/core/output/testdata/md_web_nil.golden @@ -0,0 +1,7 @@ +# Recon report · 10.0.0.1 + +Target `10.0.0.1` · Quick recon · + +--- + +No structured result was returned. diff --git a/core/output/testdata/md_web_zh.golden b/core/output/testdata/md_web_zh.golden new file mode 100644 index 00000000..d7408503 --- /dev/null +++ b/core/output/testdata/md_web_zh.golden @@ -0,0 +1,67 @@ +# 侦察报告 · 10.0.0.1 + +目标 `10.0.0.1` · 快速侦察 · + +--- + +## 概述 + +本次侦察共识别 3 台主机、4 个开放服务（含 2 个 Web 站点）。累计探测 3 条路径、命中 2 项 Web 指纹。**发现 3 项需优先复核的安全发现（凭证 / 弱口令 / 漏洞）。**另有 1 处探测报错。全程耗时 22.266s。 + +## 资产明细 + +### Example App — `http://10.0.0.1` + +- 开放服务：`tcp http 80` +- HTTP 响应：`200`、`401` +- Web 指纹：`nginx`、`basic-auth` +- 已探测路径：3 条 +- 状态：`confirmed` + +#### 分析研判 + +##### CVE-2021-41773 path traversal + +## Impact + +Arbitrary file read through `/cgi-bin/`. + +| Field | Value | +|---|---| +| CVSS | 9.8 | + +##### ssh root:root + +ssh root:root + +##### Admin console is reachable without authentication. + +##### Let me analyze the collected browser evidence. + +Let me analyze the collected browser evidence. + +## Evidence Analysis + +| Asset | Details | +|---|---| +| API | GET /api/scans | + +##### dial tcp 10.0.0.1:8443: connect: connection refused + +dial tcp 10.0.0.1:8443: connect: connection refused + +### MySQL 5.7.32 — `10.0.0.2:3306` + +- 开放服务：`tcp mysql 3306` +- 状态：`loot` + +#### 分析研判 + +##### mysql 5.7.32 + +mysql 5.7.32 + +## 其他存活主机 + +- `10.0.0.2:icmp` · icmp +- `10.0.0.3:445` · tcp smb 445 diff --git a/core/output/testdata/report_empty.json b/core/output/testdata/report_empty.json new file mode 100644 index 00000000..cad98679 --- /dev/null +++ b/core/output/testdata/report_empty.json @@ -0,0 +1,3 @@ +{ + "summary": {} +} diff --git a/core/output/testdata/report_fixture.json b/core/output/testdata/report_fixture.json new file mode 100644 index 00000000..1a4b73ba --- /dev/null +++ b/core/output/testdata/report_fixture.json @@ -0,0 +1,173 @@ +{ + "summary": { + "targets": 2, + "services": 4, + "webs": 2, + "probes": 3, + "loots": 3, + "errors": 1, + "tasks": 7, + "requests": 19, + "duration": "22.266s" + }, + "assets": [ + { + "id": "asset:http://10.0.0.1", + "key": "http://10.0.0.1", + "target": "http://10.0.0.1", + "title": "Example App", + "status": "confirmed", + "items": [ + { + "kind": "service", + "source": "gogo_portscan", + "target": "10.0.0.1:80", + "title": "http", + "summary": "nginx/1.18.0", + "tags": ["tcp", "http", "80"], + "data": {"ip": "10.0.0.1", "port": "80", "protocol": "tcp", "service": "http", "banner": "nginx/1.18.0", "is_web": true} + }, + { + "kind": "fingerprint", + "source": "gogo_portscan", + "target": "http://10.0.0.1", + "title": "nginx", + "tags": ["gogo_portscan", "nginx"], + "data": {"name": "nginx", "focus": false} + }, + { + "kind": "loot", + "source": "vuln", + "target": "http://10.0.0.1/admin", + "status": "confirmed", + "title": "CVE-2021-41773 path traversal", + "summary": "CVE-2021-41773 path traversal", + "tags": ["vuln", "apache"], + "detail": "## Impact\n\nArbitrary file read through `/cgi-bin/`.\n\n| Field | Value |\n|---|---|\n| CVSS | 9.8 |", + "data": {"kind": "vuln", "verification_status": "confirmed"} + }, + { + "kind": "loot", + "source": "weakpass", + "target": "10.0.0.1:22", + "status": "high", + "title": "ssh root:root", + "summary": "ssh root:root", + "tags": ["weakpass", "ssh"], + "data": {"kind": "weakpass"} + }, + { + "kind": "note", + "source": "deep", + "target": "http://10.0.0.1/admin", + "status": "info", + "summary": "Admin console is reachable without authentication.", + "detail": "Admin console is reachable without authentication." + }, + { + "kind": "response", + "source": "deep", + "target": "http://10.0.0.1", + "status": "response", + "detail": "Let me analyze the collected browser evidence.\n\n## Evidence Analysis\n\n| Asset | Details |\n|---|---|\n| API | GET /api/scans |" + }, + { + "kind": "path", + "source": "spray_check", + "target": "http://10.0.0.1/", + "status": "200", + "title": "Example App", + "summary": "/", + "tags": ["spray_check", "nginx", "validated"], + "data": {"url": "http://10.0.0.1/", "path": "/", "status": 200, "length": 1256, "title": "Example App", "fingers": ["nginx"], "validated": true} + }, + { + "kind": "path", + "source": "spray_check", + "target": "http://10.0.0.1/admin/login", + "status": "401", + "title": "Login", + "summary": "/admin/login", + "tags": ["spray_check", "basic-auth", "validated"], + "data": {"url": "http://10.0.0.1/admin/login", "path": "/admin/login", "status": 401, "length": 512, "title": "Login", "fingers": ["basic-auth"], "validated": true} + }, + { + "kind": "path", + "source": "spray_crawl", + "target": "http://10.0.0.1/static/app.js", + "status": "200", + "title": "js data", + "summary": "/static/app.js", + "tags": ["spray_crawl"], + "data": {"url": "http://10.0.0.1/static/app.js", "path": "/static/app.js", "status": 200, "length": 9001, "title": "js data"} + }, + { + "kind": "error", + "source": "spray_check", + "target": "scan", + "status": "error", + "summary": "dial tcp 10.0.0.1:8443: connect: connection refused", + "data": {"message": "dial tcp 10.0.0.1:8443: connect: connection refused"} + } + ] + }, + { + "id": "asset:10.0.0.2:3306", + "key": "10.0.0.2:3306", + "target": "10.0.0.2:3306", + "title": "MySQL 5.7.32", + "status": "loot", + "items": [ + { + "kind": "service", + "source": "gogo_portscan", + "target": "10.0.0.2:3306", + "title": "mysql", + "summary": "MySQL 5.7.32", + "tags": ["tcp", "mysql", "3306"], + "data": {"ip": "10.0.0.2", "port": "3306", "protocol": "tcp", "service": "mysql", "banner": "MySQL 5.7.32"} + }, + { + "kind": "loot", + "source": "fingerprint", + "target": "10.0.0.2:3306", + "status": "loot", + "title": "mysql 5.7.32", + "summary": "mysql 5.7.32", + "tags": ["fingerprint", "mysql"], + "data": {"kind": "fingerprint"} + } + ] + }, + { + "id": "asset:10.0.0.2:icmp", + "key": "10.0.0.2:icmp", + "target": "10.0.0.2:icmp", + "title": "icmp", + "items": [ + { + "kind": "service", + "source": "gogo_portscan", + "target": "10.0.0.2:icmp", + "title": "icmp", + "tags": ["icmp"], + "data": {"ip": "10.0.0.2", "protocol": "icmp", "service": "icmp"} + } + ] + }, + { + "id": "asset:10.0.0.3:445", + "key": "10.0.0.3:445", + "target": "10.0.0.3:445", + "items": [ + { + "kind": "service", + "source": "gogo_portscan", + "target": "10.0.0.3:445", + "tags": ["tcp", "smb", "445"], + "data": {"ip": "10.0.0.3", "port": "445", "protocol": "tcp", "service": "smb"} + } + ] + } + ] +} diff --git a/core/output/timeline.go b/core/output/timeline.go index 22b0ba77..da2c6b87 100644 --- a/core/output/timeline.go +++ b/core/output/timeline.go @@ -10,8 +10,8 @@ import ( "sync" "time" - "github.com/chainreactors/aiscan/pkg/aop" - xcommand "github.com/chainreactors/aiscan/pkg/aop/x/command" + "github.com/chainreactors/aiscan/core/aop" + xcommand "github.com/chainreactors/aiscan/core/aop/x/command" "github.com/chainreactors/utils/parsers" "github.com/charmbracelet/glamour" "github.com/muesli/termenv" diff --git a/core/output/timeline_test.go b/core/output/timeline_test.go index baac2462..b0d1b884 100644 --- a/core/output/timeline_test.go +++ b/core/output/timeline_test.go @@ -6,8 +6,8 @@ import ( "testing" "time" - "github.com/chainreactors/aiscan/pkg/aop" - xcommand "github.com/chainreactors/aiscan/pkg/aop/x/command" + "github.com/chainreactors/aiscan/core/aop" + xcommand "github.com/chainreactors/aiscan/core/aop/x/command" ) func TestParseLineReadsNativeAOPEnvelope(t *testing.T) { diff --git a/core/resources/keys.go b/core/resources/keys.go new file mode 100644 index 00000000..b235f9dc --- /dev/null +++ b/core/resources/keys.go @@ -0,0 +1,7 @@ +package resources + +import "github.com/chainreactors/aiscan/core/deps" + +// SetKey carries the loaded scanner resources (fingers, neutron, proton +// configs) to the command factories. +var SetKey = deps.NewKey[*Set]("core.resources.Set") diff --git a/pkg/telemetry/logger.go b/core/telemetry/logger.go similarity index 100% rename from pkg/telemetry/logger.go rename to core/telemetry/logger.go diff --git a/pkg/telemetry/logger_test.go b/core/telemetry/logger_test.go similarity index 100% rename from pkg/telemetry/logger_test.go rename to core/telemetry/logger_test.go diff --git a/pkg/telemetry/recover.go b/core/telemetry/recover.go similarity index 100% rename from pkg/telemetry/recover.go rename to core/telemetry/recover.go diff --git a/pkg/telemetry/recover_test.go b/core/telemetry/recover_test.go similarity index 100% rename from pkg/telemetry/recover_test.go rename to core/telemetry/recover_test.go diff --git a/pkg/telemetry/startup.go b/core/telemetry/startup.go similarity index 100% rename from pkg/telemetry/startup.go rename to core/telemetry/startup.go diff --git a/pkg/telemetry/startup_test.go b/core/telemetry/startup_test.go similarity index 100% rename from pkg/telemetry/startup_test.go rename to core/telemetry/startup_test.go diff --git a/pkg/agent/truncate/clip.go b/core/truncate/clip.go similarity index 100% rename from pkg/agent/truncate/clip.go rename to core/truncate/clip.go diff --git a/pkg/agent/truncate/clip_test.go b/core/truncate/clip_test.go similarity index 100% rename from pkg/agent/truncate/clip_test.go rename to core/truncate/clip_test.go diff --git a/pkg/agent/truncate/truncate.go b/core/truncate/truncate.go similarity index 99% rename from pkg/agent/truncate/truncate.go rename to core/truncate/truncate.go index 1e33e43c..a39287ca 100644 --- a/pkg/agent/truncate/truncate.go +++ b/core/truncate/truncate.go @@ -4,7 +4,7 @@ import ( "strings" "unicode/utf8" - "github.com/chainreactors/aiscan/pkg/util" + "github.com/chainreactors/aiscan/core/util" ) // ── Tier 1: 通用工具结果 (bash/read/grep/find/ls/inbox/agent result) ── diff --git a/pkg/agent/truncate/truncate_test.go b/core/truncate/truncate_test.go similarity index 100% rename from pkg/agent/truncate/truncate_test.go rename to core/truncate/truncate_test.go diff --git a/pkg/util/format.go b/core/util/format.go similarity index 100% rename from pkg/util/format.go rename to core/util/format.go diff --git a/docs/development.md b/docs/development.md index 2ed17424..06d02d51 100644 --- a/docs/development.md +++ b/docs/development.md @@ -187,7 +187,7 @@ import ( "fmt" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/core/telemetry" ) type Command struct { @@ -268,7 +268,7 @@ package whatweb import ( "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/core/telemetry" ) func init() { @@ -570,7 +570,7 @@ agent_background: true Agent 类型 Skill 通过 `skills.Store.AgentTypes()` 收集，注入到 `SubAgentTool` 中： ```go -// core/runner/runner.go +// pkg/runner/runner.go subAgentTool := agent.NewSubAgentTool(parentAgent, ib, func(name string) (agent.AgentType, error) { s, ok := rt.App.Skills.ByName(name) if !ok || !s.Agent { return error } diff --git a/docs/mechanisms.md b/docs/mechanisms.md index cd6ce37d..24b0503f 100644 --- a/docs/mechanisms.md +++ b/docs/mechanisms.md @@ -60,7 +60,7 @@ Settings UI 保存 **并发模型**: hub 的 `saveMu` 防止多个配置事务交错；本地扫描通过 managed App 租约继续使用旧运行时，不会被保存设置中断。agent 侧 `Agent.SetProvider()` / `SetMaxTurns()` 在 `mu.Lock` 下修改 `Cfg`，`Run`/`Continue` 开始时 `configSnapshot()` 在锁下拷贝，已在飞的 run 不受影响。 -**文件**: `pkg/web/service.go`, `cmd/aiscan/web_full.go`, `pkg/web/agents.go`, `pkg/webagent/agent.go`, `core/runner/runner.go`, `pkg/agent/agent.go` +**文件**: `pkg/web/service.go`, `cmd/aiscan/web_full.go`, `pkg/web/agents.go`, `pkg/webagent/agent.go`, `pkg/runner/runner.go`, `agent/agent.go` --- @@ -94,7 +94,7 @@ eval/compact 徽章仍可由 hub 从 AOP extension 派生为 Web 平台控制事 **评估器门控修正**: 旧逻辑只对 Terminated/Completed 执行评估，turn-capped（Stopped）或 token-capped（Budget）的 agent 被静默跳过。新逻辑只在 Error/Canceled 时跳过。 -**文件**: `pkg/agent/event_json.go`, `pkg/agent/evaluator/loop.go`, `pkg/web/agents.go`, `pkg/web/service.go` +**文件**: `agent/aop_emit.go`, `agent/evaluator/loop.go`, `pkg/web/agents.go`, `pkg/web/service.go` --- @@ -147,7 +147,7 @@ chat endpoint 返回 404 时包裹 actionable 建议（如"设置 `llm.provider= 这里只推断传输协议：检测 `anthropic.com` 域名选择 `anthropic`，其他自定义地址默认使用 `openai` 兼容协议。品牌默认地址由 preset 解析，不依赖域名猜测。 -**文件**: `pkg/agent/provider/anthropic.go`, `pkg/agent/provider/openai.go`, `pkg/agent/provider/http.go`, `pkg/agent/provider/provider.go` +**文件**: `agent/provider/anthropic.go`, `agent/provider/openai.go`, `agent/provider/http.go`, `agent/provider/provider.go` --- @@ -234,7 +234,7 @@ AOP error 事件把 code 保存在标准 data 中，并把 params 保存在 `ext scan、agent joined、session cleared 等产品事件保留独立的 `DomainEvent`，不携带 Agent 的 role/content/message ID 字段。 -**文件**: `core/runner/`, `pkg/aop/`, `pkg/web/service.go` +**文件**: `pkg/runner/`, `core/aop/`, `pkg/web/service.go` --- @@ -260,7 +260,7 @@ scan、agent joined、session cleared 等产品事件保留独立的 `DomainEven 跨界面 Runtime 命令通过 typed AOP command detail 标记 `presentation: preformatted`。Web timeline 在最终展示边界生成自适应 Markdown code fence；Runtime、Session 和 transport 不再处理 Markdown 或终端格式。 -**文件**: `pkg/tui/banner.go`, `pkg/tui/commands.go`, `pkg/tui/ioa.go`, `pkg/aop/x/command/command.go`, `core/output/timeline.go` +**文件**: `pkg/tui/banner.go`, `pkg/tui/commands.go`, `pkg/tui/ioa.go`, `core/aop/x/command/command.go`, `core/output/timeline.go` --- diff --git a/go.mod b/go.mod index 23da8b1b..52bb429b 100644 --- a/go.mod +++ b/go.mod @@ -6,14 +6,14 @@ require ( github.com/alecthomas/chroma/v2 v2.14.0 github.com/carapace-sh/carapace v1.11.6 github.com/chainreactors/crtm v0.0.3-0.20260618163257-073207497076 - github.com/chainreactors/fingers v1.2.2-0.20260704073236-3e22b6a528b9 - github.com/chainreactors/gogo/v2 v2.14.2-0.20260710171447-b1776cb06226 + github.com/chainreactors/fingers v1.2.2-0.20260714063144-070758342f45 + github.com/chainreactors/gogo/v2 v2.15.1-0.20260728051744-a278b33d8744 github.com/chainreactors/ioa v0.1.2-0.20260720012101-ee17a402fc18 github.com/chainreactors/libcstx/go v0.0.0-20260716111447-af6771384af7 github.com/chainreactors/logs v0.0.0-20260624034259-9aaea4aa52cc - github.com/chainreactors/neutron v0.1.1-0.20260710171341-456d36779ab2 + github.com/chainreactors/neutron v0.1.1-0.20260714062907-716c6b167cb6 github.com/chainreactors/proton v0.3.3-0.20260707162538-471f99ea6131 - github.com/chainreactors/proxyclient v1.1.1-0.20260529172347-2a80e08d5593 + github.com/chainreactors/proxyclient v1.1.1-0.20260714062913-bce898a8f790 github.com/chainreactors/proxyclient/extra v0.0.0-20260527160727-36cf133952c3 github.com/chainreactors/sdk v0.3.4-0.20260708104745-dcad8620f5e9 github.com/chainreactors/sdk/gogo v0.0.0-20260708104745-dcad8620f5e9 @@ -22,9 +22,9 @@ require ( github.com/chainreactors/spray v1.3.3-0.20260704194611-7ce7b850d447 github.com/chainreactors/tui/console v0.0.0-20260712082522-2ba36ad7841f github.com/chainreactors/tui/readline v0.0.0-20260723062039-ed89e758c21b - github.com/chainreactors/utils v0.0.0-20260707181750-8aa6ca296863 + github.com/chainreactors/utils v0.0.0-20260711153742-f3d210a5fa9d github.com/chainreactors/utils/mitmproxy v0.0.0-20260707181750-8aa6ca296863 - github.com/chainreactors/utils/parsers v0.0.3-0.20260707181750-8aa6ca296863 + github.com/chainreactors/utils/parsers v0.0.3 github.com/chainreactors/utils/pty v0.0.0-20260722180147-5b1816060721 github.com/chainreactors/zombie v1.3.0 github.com/charmbracelet/bubbles v1.0.0 diff --git a/go.sum b/go.sum index c8534106..32f93324 100644 --- a/go.sum +++ b/go.sum @@ -171,12 +171,12 @@ github.com/chainreactors/crtm v0.0.3-0.20260618163257-073207497076 h1:vIEqkeRYDy github.com/chainreactors/crtm v0.0.3-0.20260618163257-073207497076/go.mod h1:+T3JvsT0teBxi4+ValZTYWCDIwM8inbx57+nQfMFkbA= github.com/chainreactors/files v0.0.0-20240716182835-7884ee1e77f0 h1:cU3sGEODXZsUZGBXfnz0nyxF6+37vA+ZGDx6L/FKN4o= github.com/chainreactors/files v0.0.0-20240716182835-7884ee1e77f0/go.mod h1:NSxGNMRWryAyrDzZpVwmujI22wbGw6c52bQOd5zEvyU= -github.com/chainreactors/fingers v1.2.2-0.20260704073236-3e22b6a528b9 h1:6TntNzkBkKaZ2rN/w+pJUmmb0VOoZkcOZUkrIYi05RQ= -github.com/chainreactors/fingers v1.2.2-0.20260704073236-3e22b6a528b9/go.mod h1:rTZEazmD80vXmSpgwDcMo7bbZU8dop3D57XIsuy1W3M= +github.com/chainreactors/fingers v1.2.2-0.20260714063144-070758342f45 h1:wIKAvAPjQUXqAKDYG0UnpLTwn/+fwT7ipwFgNPxz88M= +github.com/chainreactors/fingers v1.2.2-0.20260714063144-070758342f45/go.mod h1:ba7u/7/I9yV7TvuWj+VV9QYz9NmlLdh6UK9kxRRft+E= github.com/chainreactors/go-re2 v1.11.1-0.20260718064805-1d8511959320 h1:gkcY9PramU2ZQZ9NWisdog4oMKaeNsd30er6Nr1pKvE= github.com/chainreactors/go-re2 v1.11.1-0.20260718064805-1d8511959320/go.mod h1:4qC68vqWSuPTct3spuTrWBqCpm00mQ707JKLS1izVjI= -github.com/chainreactors/gogo/v2 v2.14.2-0.20260710171447-b1776cb06226 h1:f/9CjNNXnSn718EsSysQUyj4AchRTw5xWqig+myOPt4= -github.com/chainreactors/gogo/v2 v2.14.2-0.20260710171447-b1776cb06226/go.mod h1:pCbDa+HwfjKCGOD4PJ0puFa/FJkE/kiy29tkX0OIw70= +github.com/chainreactors/gogo/v2 v2.15.1-0.20260728051744-a278b33d8744 h1:5Bj73ddSftvWEgjo3dEOUis1CuwPaX8JsSAIZCujfxE= +github.com/chainreactors/gogo/v2 v2.15.1-0.20260728051744-a278b33d8744/go.mod h1:Em8DiV1Rh59FCd9zN4RQBi3RnUt9yPWD/oxPEPTyXIk= github.com/chainreactors/ioa v0.1.2-0.20260720012101-ee17a402fc18 h1:X0jMNLJGBsR0prosH0sKh+ry+iyvsPifA4UlQ+iOh5M= github.com/chainreactors/ioa v0.1.2-0.20260720012101-ee17a402fc18/go.mod h1:IqHyULc67RKEmr9qsyPpJzgSGJRK8JeRXXXVthQu5Z8= github.com/chainreactors/katana v1.6.2-0.20260716115809-46dd3ac126d2 h1:pc7Vw1H4CyFnzOCxRmM1kdDwX0vJBTotIxOQjA78J/Q= @@ -185,16 +185,16 @@ github.com/chainreactors/libcstx/go v0.0.0-20260716111447-af6771384af7 h1:DJxafZ github.com/chainreactors/libcstx/go v0.0.0-20260716111447-af6771384af7/go.mod h1:YQNpUU90e8tw9k1VNEUF7smIY8kSI0J3T0/zp8ri85Y= github.com/chainreactors/logs v0.0.0-20260624034259-9aaea4aa52cc h1:e6rjnU8dmfhAnkFzLp/R5gta0LbFM13L27djxbC/i4I= github.com/chainreactors/logs v0.0.0-20260624034259-9aaea4aa52cc/go.mod h1:VrXmYPbNN5AVoo1sc5aeyPVBYqubMdb3KO/tn5rRZpo= -github.com/chainreactors/neutron v0.1.1-0.20260710171341-456d36779ab2 h1:u0gNplhf4avPGO6fowTIAl8VrXE+hKccbbrFt8ETjm4= -github.com/chainreactors/neutron v0.1.1-0.20260710171341-456d36779ab2/go.mod h1:BAWFIherRWHI1kjZkncx54tuhaKLoS6OM6T7zQBPynU= +github.com/chainreactors/neutron v0.1.1-0.20260714062907-716c6b167cb6 h1:apEDnJeZ5fe2AaEJHV6TRDldoy2t/d7E1maxmUz2tfU= +github.com/chainreactors/neutron v0.1.1-0.20260714062907-716c6b167cb6/go.mod h1:zok/CDxut71iw8NQTHKPvy0f+1G5639zr9GgDsXGnOs= github.com/chainreactors/neutron/operators/full v0.1.1-0.20260704194031-f57d0a560e32 h1:q14uQiXYcizQqkHraBghJEPzcE5StLQLIWF0HvNFKhY= github.com/chainreactors/neutron/operators/full v0.1.1-0.20260704194031-f57d0a560e32/go.mod h1:8Yg/msDYB3syXD2ryGObqSn8GG+xaBoDG/1S67A4ByQ= github.com/chainreactors/parsers v0.0.0-20260608085142-3d2c51baa8fe h1:n1pFLHHYXMiX5rCVWeciOTJUFggWXOrLtCu9jhq5Mbs= github.com/chainreactors/parsers v0.0.0-20260608085142-3d2c51baa8fe/go.mod h1:ygxMqZQ/hGY2uegUvC0LbR538hbgNH7HP4dTuv/jfSM= github.com/chainreactors/proton v0.3.3-0.20260707162538-471f99ea6131 h1:gTrBbrASTvndSBr2XL75Kdw8fAM3xw/dikTqMNzoQBE= github.com/chainreactors/proton v0.3.3-0.20260707162538-471f99ea6131/go.mod h1:c4kezBtDrE4sBIH6qF0+OShJ3/fijZENyxcUbcfZ/qQ= -github.com/chainreactors/proxyclient v1.1.1-0.20260529172347-2a80e08d5593 h1:tnXa9DobeX30xGIkiAcH+BOKKwvscvxy6GfI0Q0qu8Q= -github.com/chainreactors/proxyclient v1.1.1-0.20260529172347-2a80e08d5593/go.mod h1:xSNRChMYF8en5O5ZQVmCOmNTTyQhvsrk0D0vluh/JKk= +github.com/chainreactors/proxyclient v1.1.1-0.20260714062913-bce898a8f790 h1:We2R3rJUQBI6GGmA2kM3IOK+yvPn5NsXHhpbKiVsPE0= +github.com/chainreactors/proxyclient v1.1.1-0.20260714062913-bce898a8f790/go.mod h1:m0kmv5rgzac2q3HMYBlQULc+mubGILrPx1IiPQWcQHc= github.com/chainreactors/proxyclient/extra v0.0.0-20260527160727-36cf133952c3 h1:WVEb3Bjq3AC67oa9pNjHJfIH+mX6PozJO3S5ccVZJYQ= github.com/chainreactors/proxyclient/extra v0.0.0-20260527160727-36cf133952c3/go.mod h1:4BG8xIxTebn0VoOTEwQ2pmYIB3K2qwGINxnRHALBNZg= github.com/chainreactors/sdk v0.3.4-0.20260708104745-dcad8620f5e9 h1:zHGP9WFSpTyZYWeKDHPDoEysmNYjck25Wv8/eVtpYBE= @@ -212,14 +212,14 @@ github.com/chainreactors/tui/console v0.0.0-20260712082522-2ba36ad7841f/go.mod h github.com/chainreactors/tui/readline v0.0.0-20260723062039-ed89e758c21b h1:OeflBONN55oQ++CFJDE47pW5GfyXJpiQClFzD1aYK+o= github.com/chainreactors/tui/readline v0.0.0-20260723062039-ed89e758c21b/go.mod h1:nEHRbLD/s2GWdAGbNVjz/KDF0ac7WZ3tPMgWmW8sZWA= github.com/chainreactors/utils v0.0.0-20240716182459-e85f2b01ee16/go.mod h1:LajXuvESQwP+qCMAvlcoSXppQCjuLlBrnQpu9XQ1HtU= -github.com/chainreactors/utils v0.0.0-20260707181750-8aa6ca296863 h1:jfuZD+vg3/K/+l8au9RXnZCQf0J6S3vG6VRqmeBmxo0= -github.com/chainreactors/utils v0.0.0-20260707181750-8aa6ca296863/go.mod h1:xwbUlFoSSxLHujyb8D48o1s2DqmEAxUNfxIy0DVUmcg= +github.com/chainreactors/utils v0.0.0-20260711153742-f3d210a5fa9d h1:wlJ6oMbVLKrpxHmaXGSxmJt1F8l3kvqily0N58FGfLM= +github.com/chainreactors/utils v0.0.0-20260711153742-f3d210a5fa9d/go.mod h1:xwbUlFoSSxLHujyb8D48o1s2DqmEAxUNfxIy0DVUmcg= github.com/chainreactors/utils/cert v0.0.0-20260707181750-8aa6ca296863 h1:41tvJzi9t1NUlM/CzVdl8OG+W6PMFChsDOChohI2VeU= github.com/chainreactors/utils/cert v0.0.0-20260707181750-8aa6ca296863/go.mod h1:xvvWMcU9Fcht6GR1cc9ceAZ3/Hl2HrkoRzpeyOzx1rQ= github.com/chainreactors/utils/mitmproxy v0.0.0-20260707181750-8aa6ca296863 h1:r6UUUUQt4r/0SL6vgrwoq6ynidAkN3auSZsvzZ5BBRE= github.com/chainreactors/utils/mitmproxy v0.0.0-20260707181750-8aa6ca296863/go.mod h1:2O3/Vw66VnbzhwsHGFJ2Ge98RuSh6XzMMFGZmMmlZ9M= -github.com/chainreactors/utils/parsers v0.0.3-0.20260707181750-8aa6ca296863 h1:u9cXebLoVtKwN0KkpGFfrjANUYo+93MijB//X9qONeY= -github.com/chainreactors/utils/parsers v0.0.3-0.20260707181750-8aa6ca296863/go.mod h1:S9lkpQ1I4wcBq0YEBde/UPmR061IPok3bLl7aPz6Vkk= +github.com/chainreactors/utils/parsers v0.0.3 h1:3ld7xG5TSvzikVOCkQHjqjHO3otjODwSwHQDkMKbu5o= +github.com/chainreactors/utils/parsers v0.0.3/go.mod h1:bE/znJWt08n9QOORWsWu0ggB8GWfOg3+dfUMMITmwV4= github.com/chainreactors/utils/pty v0.0.0-20260722180147-5b1816060721 h1:gxkedbTvFEFTtel7XJEPMVh1iznfD+91woPkGBXZMNk= github.com/chainreactors/utils/pty v0.0.0-20260722180147-5b1816060721/go.mod h1:RW1v+8hFMeO9+TJyQ1iIx9Ea37s+B7BaDf9fyJ0OEC4= github.com/chainreactors/words v0.0.0-20260520145736-270600e60fb4 h1:lvnDYEkatmZFHP5i321qQXK9L4vKRfso/uUfr5tOeC8= diff --git a/pkg/commands/bash.go b/pkg/commands/bash.go index f50aca9b..720832eb 100644 --- a/pkg/commands/bash.go +++ b/pkg/commands/bash.go @@ -10,10 +10,10 @@ import ( "strings" "time" + "github.com/chainreactors/aiscan/agent/inbox" + "github.com/chainreactors/aiscan/agent/tmux" coretool "github.com/chainreactors/aiscan/core/tool" - "github.com/chainreactors/aiscan/pkg/agent/inbox" - "github.com/chainreactors/aiscan/pkg/agent/tmux" - "github.com/chainreactors/aiscan/pkg/agent/truncate" + "github.com/chainreactors/aiscan/core/truncate" ) const ( @@ -110,7 +110,7 @@ func (t *BashTool) Execute(ctx context.Context, arguments string) (coretool.Resu return coretool.Result{}, err } - return t.waitOrBackground(execution, ctx, inboxFromContext(ctx)), nil + return t.waitOrBackground(execution, ctx, inbox.FromContext(ctx)), nil } // RunForeground executes command through the same tmux/registered-command diff --git a/pkg/commands/bash_inbox_test.go b/pkg/commands/bash_inbox_test.go index 4441bcb4..c3024f0d 100644 --- a/pkg/commands/bash_inbox_test.go +++ b/pkg/commands/bash_inbox_test.go @@ -6,7 +6,7 @@ import ( "testing" "time" - "github.com/chainreactors/aiscan/pkg/agent/inbox" + "github.com/chainreactors/aiscan/agent/inbox" ) func TestBashBackgroundMonitorUsesInvocationInbox(t *testing.T) { diff --git a/pkg/commands/bash_test.go b/pkg/commands/bash_test.go index cdd41c81..b6a2c589 100644 --- a/pkg/commands/bash_test.go +++ b/pkg/commands/bash_test.go @@ -14,10 +14,10 @@ import ( "testing" "time" + tmux "github.com/chainreactors/aiscan/agent/tmux" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/core/tool" - tmux "github.com/chainreactors/aiscan/pkg/agent/tmux" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" ) // --------------------------------------------------------------------------- diff --git a/pkg/commands/command.go b/pkg/commands/command.go index bd5fbc8f..2e1775a0 100644 --- a/pkg/commands/command.go +++ b/pkg/commands/command.go @@ -7,8 +7,8 @@ import ( "strings" "sync" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/core/tool" - "github.com/chainreactors/aiscan/pkg/telemetry" ) var _ tool.Executor = (*CommandRegistry)(nil) diff --git a/pkg/commands/context.go b/pkg/commands/context.go deleted file mode 100644 index 70d44c52..00000000 --- a/pkg/commands/context.go +++ /dev/null @@ -1,27 +0,0 @@ -package commands - -import ( - "context" - - "github.com/chainreactors/aiscan/pkg/agent/inbox" -) - -type inboxContextKey struct{} - -// ContextWithInbox scopes asynchronous command notifications to the agent -// session that invoked the tool. -func ContextWithInbox(ctx context.Context, ib inbox.Inbox) context.Context { - if ctx == nil { - ctx = context.Background() - } - return context.WithValue(ctx, inboxContextKey{}, ib) -} - -func inboxFromContext(ctx context.Context) inbox.Inbox { - if ctx != nil { - if ib, ok := ctx.Value(inboxContextKey{}).(inbox.Inbox); ok && ib != nil { - return ib - } - } - return nil -} diff --git a/pkg/commands/execution.go b/pkg/commands/execution.go index 67ee0ea1..699a4a2a 100644 --- a/pkg/commands/execution.go +++ b/pkg/commands/execution.go @@ -6,7 +6,7 @@ import ( "sync" "time" - "github.com/chainreactors/aiscan/pkg/agent/tmux" + "github.com/chainreactors/aiscan/agent/tmux" ) // Execution is one shell or built-in command invocation. Its ID is always the diff --git a/pkg/commands/factory.go b/pkg/commands/factory.go index 9847a050..3ba9b0c9 100644 --- a/pkg/commands/factory.go +++ b/pkg/commands/factory.go @@ -3,33 +3,58 @@ package commands import ( "sync" + "github.com/chainreactors/aiscan/agent/hooks" + "github.com/chainreactors/aiscan/agent/provider" + "github.com/chainreactors/aiscan/core/capability" + "github.com/chainreactors/aiscan/core/deps" "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/core/telemetry" ) type Factory struct { - Group string - Build func(deps *Deps, reg *CommandRegistry) + Capability capability.ID + Build func(deps *Deps, reg *CommandRegistry) } +// SkillSource is the slice of skills.Store the built-in tools need; declaring +// it here keeps pkg/commands from importing the skill store itself. +type SkillSource interface { + VirtualFileReader + VirtualGlobber +} + +// Deps carries everything a Factory may need. Values whose type pkg/commands +// must not import (scanner engines, resources, IOA client, scan options) travel +// in the Bag under keys owned by their own package. type Deps struct { + *deps.Bag + WorkDir string BashTimeout int - SkillStore any + SkillStore SkillSource RunnerMode bool - EngineSet any - Resources any - IOAClient any - Provider any + Provider provider.Provider ScannerProxy string - ScanOpts []any Logger telemetry.Logger NodeName string NodeMeta map[string]any TavilyKeys string // comma-separated Tavily API keys (build-time fallback) DataBus *eventbus.Bus[output.ToolDataEvent] + Hooks *hooks.Registry +} + +// Provide stores a typed dependency, allocating the bag on first use so a +// literal-constructed Deps cannot drop it silently. +func Provide[T any](d *Deps, key deps.Key[T], value T) { + if d == nil { + return + } + if d.Bag == nil { + d.Bag = deps.New() + } + deps.Set(d.Bag, key, value) } func (d *Deps) GetLogger() telemetry.Logger { @@ -39,6 +64,13 @@ func (d *Deps) GetLogger() telemetry.Logger { return telemetry.NopLogger() } +// Skip reports that a factory bailed out because a dependency is missing. It +// exists because these sites used to return silently, leaving the tool absent +// from the registry with nothing in the log to explain it. +func (d *Deps) Skip(id, dep string) { + d.GetLogger().Warnf("%s", telemetry.StartupLine("skip", id, "missing dependency "+dep)) +} + var ( factoryMu sync.Mutex factories []Factory @@ -50,25 +82,16 @@ func RegisterFactory(f Factory) { factories = append(factories, f) } -func BuildAll(deps *Deps, reg *CommandRegistry) { - factoryMu.Lock() - snapshot := make([]Factory, len(factories)) - copy(snapshot, factories) - factoryMu.Unlock() - - for _, f := range snapshot { - f.Build(deps, reg) - } -} - -func BuildGroup(group string, deps *Deps, reg *CommandRegistry) { +// BuildPlan is the only factory assembly path. Factory membership is keyed by +// capability identity; groups are selection metadata owned by descriptors. +func BuildPlan(plan capability.Plan, deps *Deps, reg *CommandRegistry) { factoryMu.Lock() snapshot := make([]Factory, len(factories)) copy(snapshot, factories) factoryMu.Unlock() for _, f := range snapshot { - if f.Group != group { + if f.Capability == "" || !plan.Has(f.Capability) { continue } f.Build(deps, reg) diff --git a/pkg/commands/glob.go b/pkg/commands/glob.go index 1006ffa5..c86a75f7 100644 --- a/pkg/commands/glob.go +++ b/pkg/commands/glob.go @@ -9,7 +9,7 @@ import ( "strings" coretool "github.com/chainreactors/aiscan/core/tool" - "github.com/chainreactors/aiscan/pkg/agent/truncate" + "github.com/chainreactors/aiscan/core/truncate" ) const maxGlobResults = truncate.MaxGlobResults diff --git a/pkg/commands/image_optimize.go b/pkg/commands/image_optimize.go index 4295b9b8..f7d21e96 100644 --- a/pkg/commands/image_optimize.go +++ b/pkg/commands/image_optimize.go @@ -14,7 +14,7 @@ import ( ) const ( - maxDimension = 2000 + maxDimension = 2000 maxPayloadBytes = 4_500_000 // 4.5MB base64, below Anthropic's 5MB limit ) diff --git a/pkg/commands/list.go b/pkg/commands/list.go index 62b4da28..19007547 100644 --- a/pkg/commands/list.go +++ b/pkg/commands/list.go @@ -8,7 +8,7 @@ import ( "path/filepath" coretool "github.com/chainreactors/aiscan/core/tool" - "github.com/chainreactors/aiscan/pkg/agent/truncate" + "github.com/chainreactors/aiscan/core/truncate" ) // ListTool lists directory entries through the host filesystem API. diff --git a/pkg/commands/read.go b/pkg/commands/read.go index 29e0984e..7d1cc580 100644 --- a/pkg/commands/read.go +++ b/pkg/commands/read.go @@ -10,7 +10,7 @@ import ( "unicode/utf8" coretool "github.com/chainreactors/aiscan/core/tool" - "github.com/chainreactors/aiscan/pkg/agent/truncate" + "github.com/chainreactors/aiscan/core/truncate" ) const ( diff --git a/pkg/commands/register.go b/pkg/commands/register.go index 0efd9294..367fddc2 100644 --- a/pkg/commands/register.go +++ b/pkg/commands/register.go @@ -1,11 +1,15 @@ package commands +import "github.com/chainreactors/aiscan/core/capability" + func init() { + capability.Register(capability.Descriptor{ID: "core", Kind: capability.KindTool, Group: "core"}) RegisterFactory(Factory{ - Group: "core", + Capability: "core", Build: func(deps *Deps, reg *CommandRegistry) { workDir := deps.WorkDir if workDir == "" { + deps.Skip("core", "WorkDir") return } timeout := deps.BashTimeout @@ -15,12 +19,8 @@ func init() { var readers []VirtualFileReader var globbers []VirtualGlobber if deps.SkillStore != nil { - if r, ok := deps.SkillStore.(VirtualFileReader); ok { - readers = append(readers, r) - } - if g, ok := deps.SkillStore.(VirtualGlobber); ok { - globbers = append(globbers, g) - } + readers = append(readers, deps.SkillStore) + globbers = append(globbers, deps.SkillStore) } reg.RegisterTool(NewReadTool(workDir, readers...)) reg.RegisterTool(NewWriteTool(workDir)) diff --git a/pkg/commands/register_test.go b/pkg/commands/register_test.go index 072bc4eb..1b625953 100644 --- a/pkg/commands/register_test.go +++ b/pkg/commands/register_test.go @@ -1,6 +1,10 @@ package commands -import "testing" +import ( + "testing" + + "github.com/chainreactors/aiscan/core/capability" +) func closeRegistryTools(registry *CommandRegistry) { for _, tool := range registry.Tools() { @@ -12,14 +16,14 @@ func closeRegistryTools(registry *CommandRegistry) { func TestNativeListToolIsRunnerOnly(t *testing.T) { regular := NewRegistry() - BuildGroup("core", &Deps{WorkDir: t.TempDir()}, regular) + BuildPlan(capability.Select(capability.Options{Groups: []string{"core"}}), &Deps{WorkDir: t.TempDir()}, regular) defer closeRegistryTools(regular) if _, ok := regular.GetTool("ls"); ok { t.Fatal("regular agent must not expose the runner-only ls tool") } runner := NewRegistry() - BuildGroup("core", &Deps{WorkDir: t.TempDir(), RunnerMode: true}, runner) + BuildPlan(capability.Select(capability.Options{Groups: []string{"core"}}), &Deps{WorkDir: t.TempDir(), RunnerMode: true}, runner) defer closeRegistryTools(runner) if _, ok := runner.GetTool("ls"); !ok { t.Fatal("runner mode must expose the native ls tool") diff --git a/pkg/commands/tmux.go b/pkg/commands/tmux.go index ff0ad471..a22a40ac 100644 --- a/pkg/commands/tmux.go +++ b/pkg/commands/tmux.go @@ -7,8 +7,8 @@ import ( "strings" "time" - "github.com/chainreactors/aiscan/pkg/agent/tmux" - "github.com/chainreactors/aiscan/pkg/agent/truncate" + "github.com/chainreactors/aiscan/agent/tmux" + "github.com/chainreactors/aiscan/core/truncate" ) type tmuxCommand struct { diff --git a/pkg/commands/tmux_test.go b/pkg/commands/tmux_test.go index d20841c8..3c32c0dc 100644 --- a/pkg/commands/tmux_test.go +++ b/pkg/commands/tmux_test.go @@ -9,7 +9,7 @@ import ( "testing" "time" - tmuxpkg "github.com/chainreactors/aiscan/pkg/agent/tmux" + tmuxpkg "github.com/chainreactors/aiscan/agent/tmux" ) type testOutputWriter struct{ bytes.Buffer } diff --git a/pkg/commands/write.go b/pkg/commands/write.go index e30c09a8..1554b3b4 100644 --- a/pkg/commands/write.go +++ b/pkg/commands/write.go @@ -9,7 +9,7 @@ import ( "strings" coretool "github.com/chainreactors/aiscan/core/tool" - "github.com/chainreactors/aiscan/pkg/agent/truncate" + "github.com/chainreactors/aiscan/core/truncate" ) type WriteTool struct { diff --git a/pkg/agent/probe/config.go b/pkg/probe/config.go similarity index 100% rename from pkg/agent/probe/config.go rename to pkg/probe/config.go diff --git a/pkg/agent/probe/conn.go b/pkg/probe/conn.go similarity index 100% rename from pkg/agent/probe/conn.go rename to pkg/probe/conn.go diff --git a/core/runner/app.go b/pkg/runner/app.go similarity index 89% rename from core/runner/app.go rename to pkg/runner/app.go index ee718b8c..95e0d8ef 100644 --- a/core/runner/app.go +++ b/pkg/runner/app.go @@ -8,16 +8,19 @@ import ( "sync" "time" - cfg "github.com/chainreactors/aiscan/core/config" + "github.com/chainreactors/aiscan/agent" + "github.com/chainreactors/aiscan/agent/hooks" + "github.com/chainreactors/aiscan/agent/probe" + "github.com/chainreactors/aiscan/core/capability" "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" - "github.com/chainreactors/aiscan/pkg/agent" - "github.com/chainreactors/aiscan/pkg/agent/probe" - "github.com/chainreactors/aiscan/pkg/agent/truncate" + "github.com/chainreactors/aiscan/core/telemetry" + "github.com/chainreactors/aiscan/core/truncate" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" "github.com/chainreactors/aiscan/skills" + ioatools "github.com/chainreactors/aiscan/tools/ioa" ioaclient "github.com/chainreactors/ioa/client" + "github.com/chainreactors/ioa/protocols" ) type App struct { @@ -25,6 +28,7 @@ type App struct { ProviderConfig agent.ProviderConfig ProviderFallbacks []agent.ProviderEntry Commands *commands.CommandRegistry + Hooks *hooks.Registry Engines any Skills *skills.Store SkillDiagnostics []skills.Diagnostic @@ -37,7 +41,7 @@ type App struct { logger telemetry.Logger } -func NewApp(ctx context.Context, rc cfg.RuntimeConfig) (*App, error) { +func NewApp(ctx context.Context, rc ApplicationConfig) (*App, error) { a := &App{} logger := rc.Logger if logger == nil { @@ -45,6 +49,10 @@ func NewApp(ctx context.Context, rc cfg.RuntimeConfig) (*App, error) { } a.logger = logger logger = a.Logger() + a.Hooks = hooks.New() + a.Hooks.SetErrorSink(func(he *hooks.HandlerError) { + a.Logger().Warnf("hook failed kind=%s source=%s error=%q", he.Kind, he.Source, he.Err) + }) a.DataBus = eventbus.New[output.ToolDataEvent]() a.SCOSidecar = output.NewSCOSidecar(a.DataBus, output.CSTXTransform) @@ -79,7 +87,7 @@ func NewApp(ctx context.Context, rc cfg.RuntimeConfig) (*App, error) { } } - a.Commands = initCoreCommands(rc, a.Provider, a.Skills, logger) + a.Commands = initCoreCommands(rc, a.Provider, a.Skills, a.Hooks, logger) a.enginesReady = make(chan struct{}) go func() { @@ -232,11 +240,7 @@ func llmConfigLabel(providerName, model string) string { return providerName + "/" + model } -// optionalToolGroups lists all selectable tool groups that can be enabled via -// --tools or config. Arsenal is always loaded and is NOT in this list. -var optionalToolGroups = []string{"search", "browser"} - -func initCoreCommands(rc cfg.RuntimeConfig, llmProvider agent.Provider, skillStore *skills.Store, logger telemetry.Logger) *commands.CommandRegistry { +func initCoreCommands(rc ApplicationConfig, llmProvider agent.Provider, skillStore *skills.Store, hookRegistry *hooks.Registry, logger telemetry.Logger) *commands.CommandRegistry { cmdReg := commands.NewRegistry() workDir, _ := os.Getwd() deps := &commands.Deps{ @@ -246,20 +250,13 @@ func initCoreCommands(rc cfg.RuntimeConfig, llmProvider agent.Provider, skillSto Provider: llmProvider, Logger: logger, TavilyKeys: rc.Tools.TavilyKeys, + Hooks: hookRegistry, } - commands.BuildGroup("core", deps, cmdReg) - commands.BuildGroup("arsenal", deps, cmdReg) - - enabled := rc.Tools.OptionalTools - if len(enabled) == 0 { - for _, g := range optionalToolGroups { - commands.BuildGroup(g, deps, cmdReg) - } - } else { - for _, g := range enabled { - commands.BuildGroup(g, deps, cmdReg) - } - } + plan := capability.Select(capability.Options{ + Groups: []string{"core", "arsenal", "search", "browser"}, + OptionalTools: rc.Tools.OptionalTools, + }) + commands.BuildPlan(plan, deps, cmdReg) return cmdReg } @@ -321,7 +318,7 @@ func quoteCommandArg(value string) string { return `"` + value + `"` } -func (a *App) InitIOA(ctx context.Context, ioa cfg.IOAConfig) error { +func (a *App) InitIOA(ctx context.Context, ioa IOAConfig) error { client, err := newIOAClient(ioa) if err != nil { return err @@ -338,11 +335,11 @@ func (a *App) InitIOA(ctx context.Context, ioa cfg.IOAConfig) error { a.IOAStreamClient = client if ioa.RegisterTools && a.Commands != nil { deps := &commands.Deps{ - IOAClient: client, - NodeName: ioa.NodeName, - NodeMeta: ioa.NodeMeta, + NodeName: ioa.NodeName, + NodeMeta: ioa.NodeMeta, } - commands.BuildGroup("ioa", deps, a.Commands) + commands.Provide(deps, ioatools.ClientKey, protocols.ClientAPI(client)) + commands.BuildPlan(capability.Select(capability.Options{Groups: []string{"ioa"}}), deps, a.Commands) } if ioa.AutoRegister { if err := client.EnsureRegistered(ctx, ioa.NodeName, "", ioa.NodeMeta); err != nil { @@ -355,7 +352,7 @@ func (a *App) InitIOA(ctx context.Context, ioa cfg.IOAConfig) error { return nil } -func (a *App) retryIOARegistration(ctx context.Context, client *ioaclient.Client, ioa cfg.IOAConfig) { +func (a *App) retryIOARegistration(ctx context.Context, client *ioaclient.Client, ioa IOAConfig) { for attempt := 0; ; attempt++ { delay := agent.RetryDelay(attempt) select { @@ -371,7 +368,7 @@ func (a *App) retryIOARegistration(ctx context.Context, client *ioaclient.Client } } -func (a *App) configureIOASpace(ctx context.Context, client *ioaclient.Client, ioa cfg.IOAConfig) { +func (a *App) configureIOASpace(ctx context.Context, client *ioaclient.Client, ioa IOAConfig) { if ioa.Space != "" && client != nil && client.Bound() { info, err := client.Space(ctx, ioa.Space, "aiscan agent") if err == nil { @@ -388,7 +385,7 @@ func (a *App) setIOASpace(spaceID string) { } } -func newIOAClient(ioa cfg.IOAConfig) (*ioaclient.Client, error) { +func newIOAClient(ioa IOAConfig) (*ioaclient.Client, error) { if ioa.URL == "" { return nil, nil } diff --git a/core/runner/app_test.go b/pkg/runner/app_test.go similarity index 96% rename from core/runner/app_test.go rename to pkg/runner/app_test.go index 86cc4343..daa7335b 100644 --- a/core/runner/app_test.go +++ b/pkg/runner/app_test.go @@ -9,8 +9,8 @@ import ( "strings" "testing" - "github.com/chainreactors/aiscan/pkg/agent" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/agent" + "github.com/chainreactors/aiscan/core/telemetry" ) func TestLogLLMProbeStatusReady(t *testing.T) { diff --git a/core/config/app_config.go b/pkg/runner/application_builder.go similarity index 74% rename from core/config/app_config.go rename to pkg/runner/application_builder.go index 7f21298e..48562905 100644 --- a/core/config/app_config.go +++ b/pkg/runner/application_builder.go @@ -1,9 +1,10 @@ -package config +package runner import ( "strings" - "github.com/chainreactors/aiscan/pkg/telemetry" + cfg "github.com/chainreactors/aiscan/core/config" + "github.com/chainreactors/aiscan/core/telemetry" ) type RuntimeFeatures struct { @@ -15,9 +16,9 @@ type RuntimeFeatures struct { Warning string } -func AppConfig(option *Option, features RuntimeFeatures, logger telemetry.Logger) RuntimeConfig { - return RuntimeConfig{ - Provider: RuntimeProviderConfig{ +func AppConfig(option *cfg.Option, features RuntimeFeatures, logger telemetry.Logger) ApplicationConfig { + return ApplicationConfig{ + Provider: ApplicationProviderConfig{ Enabled: features.ProviderEnabled, Config: ProviderConfig(option), Fallbacks: FallbackProviderConfigs(option), @@ -28,7 +29,7 @@ func AppConfig(option *Option, features RuntimeFeatures, logger telemetry.Logger CyberhubKey: option.CyberhubKey, CyberhubMode: option.CyberhubMode, AIEnabled: features.AIEnabled, - VerifyMode: ResolveString(option.ScanConfig.Verify, DefaultVerify), + VerifyMode: cfg.ResolveString(option.ScanConfig.Verify, cfg.DefaultVerify), Proxy: option.Proxy, FofaEmail: option.FofaEmail, FofaKey: option.FofaKey, @@ -40,7 +41,7 @@ func AppConfig(option *Option, features RuntimeFeatures, logger telemetry.Logger Tools: ToolConfig{ Enabled: features.ToolsEnabled, BashTimeout: 300, - TavilyKeys: resolveTavilyKeys(option.TavilyKey, ResolveString(option.SearchConfig.TavilyKeys, DefaultTavilyKeys)), + TavilyKeys: resolveTavilyKeys(option.TavilyKey, cfg.ResolveString(option.SearchConfig.TavilyKeys, cfg.DefaultTavilyKeys)), OptionalTools: option.Tools, }, Logger: logger, @@ -48,7 +49,7 @@ func AppConfig(option *Option, features RuntimeFeatures, logger telemetry.Logger } } -func skillPathsFromOptions(option *Option) []string { +func skillPathsFromOptions(option *cfg.Option) []string { var paths []string for _, s := range option.Skills { if looksLikePath(s) { diff --git a/core/config/runtime.go b/pkg/runner/application_config.go similarity index 81% rename from core/config/runtime.go rename to pkg/runner/application_config.go index 3ee27c91..8aaade0c 100644 --- a/core/config/runtime.go +++ b/pkg/runner/application_config.go @@ -1,13 +1,13 @@ -package config +package runner import ( - "github.com/chainreactors/aiscan/pkg/agent" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/agent" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/ioa/protocols" ) -type RuntimeConfig struct { - Provider RuntimeProviderConfig +type ApplicationConfig struct { + Provider ApplicationProviderConfig Scanner ScannerConfig Tools ToolConfig IOA *IOAConfig @@ -16,7 +16,7 @@ type RuntimeConfig struct { SkipEngines bool } -type RuntimeProviderConfig struct { +type ApplicationProviderConfig struct { Enabled bool Config agent.ProviderConfig Fallbacks []agent.ProviderConfig diff --git a/core/runner/hooks.go b/pkg/runner/hooks.go similarity index 86% rename from core/runner/hooks.go rename to pkg/runner/hooks.go index 240608a1..55ff5021 100644 --- a/core/runner/hooks.go +++ b/pkg/runner/hooks.go @@ -4,12 +4,12 @@ import ( "context" cfg "github.com/chainreactors/aiscan/core/config" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/core/telemetry" ) // ScannerInitFunc initializes scanner engines and registers scanner commands. // Set via init() from the package imported by cmd/aiscan. -var ScannerInitFunc func(ctx context.Context, a *App, rc cfg.RuntimeConfig, logger telemetry.Logger) +var ScannerInitFunc func(ctx context.Context, a *App, rc ApplicationConfig, logger telemetry.Logger) // ScannerWithAgentFunc runs a scanner command with AI agent assistance. // Set via init() from the package imported by cmd/aiscan. diff --git a/core/runner/ioa.go b/pkg/runner/ioa.go similarity index 95% rename from core/runner/ioa.go rename to pkg/runner/ioa.go index ddf132ff..121edf9b 100644 --- a/core/runner/ioa.go +++ b/pkg/runner/ioa.go @@ -9,7 +9,7 @@ import ( "time" cfg "github.com/chainreactors/aiscan/core/config" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/core/telemetry" ) func RunIOAServe(ctx context.Context, option *cfg.Option, logger telemetry.Logger) error { diff --git a/core/runner/local_repl.go b/pkg/runner/local_repl.go similarity index 100% rename from core/runner/local_repl.go rename to pkg/runner/local_repl.go diff --git a/pkg/agent/loop_tool.go b/pkg/runner/loop_command.go similarity index 70% rename from pkg/agent/loop_tool.go rename to pkg/runner/loop_command.go index e059b49c..14617a4d 100644 --- a/pkg/agent/loop_tool.go +++ b/pkg/runner/loop_command.go @@ -1,4 +1,4 @@ -package agent +package runner import ( "context" @@ -7,6 +7,7 @@ import ( "strings" "time" + "github.com/chainreactors/aiscan/agent" "github.com/chainreactors/aiscan/pkg/commands" ) @@ -16,32 +17,13 @@ import ( // bash(command="loop 30s check status") // bash(command="loop list") // bash(command="loop stop loop-a1b2c3d4") -type LoopCommand struct{} +type loopCommand struct{} -type loopSchedulerContextKey struct{} +func newLoopCommand() *loopCommand { return &loopCommand{} } -// ContextWithLoopScheduler scopes direct REPL command execution to one runtime -// session. Agent tool calls obtain the scheduler from their Config snapshot. -func ContextWithLoopScheduler(ctx context.Context, scheduler *LoopScheduler) context.Context { - if ctx == nil { - ctx = context.Background() - } - return context.WithValue(ctx, loopSchedulerContextKey{}, scheduler) -} - -func loopSchedulerFromContext(ctx context.Context) *LoopScheduler { - if ctx == nil { - return nil - } - scheduler, _ := ctx.Value(loopSchedulerContextKey{}).(*LoopScheduler) - return scheduler -} +func (c *loopCommand) Name() string { return "loop" } -func NewLoopCommand() *LoopCommand { return &LoopCommand{} } - -func (c *LoopCommand) Name() string { return "loop" } - -func (c *LoopCommand) Usage() string { +func (c *loopCommand) Usage() string { return `loop — recurring task scheduler Usage: @@ -62,13 +44,8 @@ Examples: loop 5m monitor targets every 5 minutes` } -func (c *LoopCommand) Run(ctx context.Context, execution *commands.Execution) (any, error) { - scheduler := loopSchedulerFromContext(ctx) - if scheduler == nil { - if cfg, ok := toolAgentConfig(ctx); ok { - scheduler = cfg.LoopScheduler - } - } +func (c *loopCommand) Run(ctx context.Context, execution *commands.Execution) (any, error) { + scheduler := agent.LoopSchedulerFromContext(ctx) if scheduler == nil { return nil, fmt.Errorf("loop scheduler is not configured") } @@ -99,12 +76,12 @@ func (c *LoopCommand) Run(ctx context.Context, execution *commands.Execution) (a } } -func (c *LoopCommand) create(ctx context.Context, scheduler *LoopScheduler, output io.Writer, args []string) error { +func (c *loopCommand) create(ctx context.Context, scheduler *agent.LoopScheduler, output io.Writer, args []string) error { if len(args) < 2 { return fmt.Errorf("usage: loop ") } - entry := LoopEntry{Mode: ModeInbox} + entry := agent.LoopEntry{Mode: agent.ModeInbox} // try cron first (5 space-separated fields), then duration if cron, rest, ok := tryCronPrefix(args); ok { @@ -131,9 +108,9 @@ func (c *LoopCommand) create(ctx context.Context, scheduler *LoopScheduler, outp // tryCronPrefix attempts to parse the first 5 args as a cron expression. // Returns the parsed expression, remaining args, and whether it succeeded. -func tryCronPrefix(args []string) (*CronExpr, []string, bool) { +func tryCronPrefix(args []string) (*agent.CronExpr, []string, bool) { if len(args) >= 2 && strings.Contains(args[0], " ") { - if cron, err := ParseCron(args[0]); err == nil { + if cron, err := agent.ParseCron(args[0]); err == nil { return cron, args[1:], true } } @@ -141,14 +118,14 @@ func tryCronPrefix(args []string) (*CronExpr, []string, bool) { return nil, nil, false } expr := strings.Join(args[:5], " ") - cron, err := ParseCron(expr) + cron, err := agent.ParseCron(expr) if err != nil { return nil, nil, false } return cron, args[5:], true } -func (c *LoopCommand) list(scheduler *LoopScheduler, output io.Writer) error { +func (c *loopCommand) list(scheduler *agent.LoopScheduler, output io.Writer) error { loops := scheduler.List() if len(loops) == 0 { _, _ = fmt.Fprint(output, "No active loops.\n") @@ -167,7 +144,7 @@ func (c *LoopCommand) list(scheduler *LoopScheduler, output io.Writer) error { return nil } -func (c *LoopCommand) stop(scheduler *LoopScheduler, output io.Writer, name string) error { +func (c *loopCommand) stop(scheduler *agent.LoopScheduler, output io.Writer, name string) error { if err := scheduler.Remove(name); err != nil { return err } diff --git a/core/runner/prompt.go b/pkg/runner/prompt.go similarity index 99% rename from core/runner/prompt.go rename to pkg/runner/prompt.go index daf5e653..d3d42958 100644 --- a/core/runner/prompt.go +++ b/pkg/runner/prompt.go @@ -7,7 +7,7 @@ import ( "text/template" "time" - "github.com/chainreactors/aiscan/pkg/agent" + "github.com/chainreactors/aiscan/agent" "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/aiscan/skills" ) diff --git a/core/runner/prompt_test.go b/pkg/runner/prompt_test.go similarity index 98% rename from core/runner/prompt_test.go rename to pkg/runner/prompt_test.go index 3970a9d0..fc32aae1 100644 --- a/core/runner/prompt_test.go +++ b/pkg/runner/prompt_test.go @@ -6,8 +6,8 @@ import ( "testing" cfg "github.com/chainreactors/aiscan/core/config" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" "github.com/chainreactors/aiscan/skills" ) diff --git a/core/config/provider.go b/pkg/runner/provider_config.go similarity index 62% rename from core/config/provider.go rename to pkg/runner/provider_config.go index f31efdca..de1eff24 100644 --- a/core/config/provider.go +++ b/pkg/runner/provider_config.go @@ -1,43 +1,24 @@ -package config +package runner -import "github.com/chainreactors/aiscan/pkg/agent" - -var ( - DefaultProvider = "openai" - DefaultBaseURL = "" - DefaultAPIKey = "" - DefaultModel = "" - - DefaultScannerProxy = "" - - DefaultCyberhubURL = "" - DefaultCyberhubKey = "" - DefaultCyberhubMode = "merge" - - DefaultVerify = "auto" - - DefaultIOAURL = "" - DefaultIOANodeID = "" - DefaultIOANodeName = "" - DefaultSpace = "" - - DefaultTavilyKeys = "" +import ( + "github.com/chainreactors/aiscan/agent" + cfg "github.com/chainreactors/aiscan/core/config" ) func defaultProviderConfig() agent.ProviderConfig { return agent.ProviderConfig{ - Provider: DefaultProvider, - BaseURL: DefaultBaseURL, - APIKey: DefaultAPIKey, - Model: DefaultModel, + Provider: cfg.DefaultProvider, + BaseURL: cfg.DefaultBaseURL, + APIKey: cfg.DefaultAPIKey, + Model: cfg.DefaultModel, } } -func hasSingleProviderFields(option *Option) bool { +func hasSingleProviderFields(option *cfg.Option) bool { return option.Provider != "" || option.BaseURL != "" || option.APIKey != "" || option.Model != "" } -func entryToProviderConfig(entry LLMProviderEntry) agent.ProviderConfig { +func entryToProviderConfig(entry cfg.LLMProviderEntry) agent.ProviderConfig { cfg := agent.ProviderConfig{ Provider: entry.Provider, BaseURL: entry.BaseURL, @@ -57,7 +38,7 @@ func entryToProviderConfig(entry LLMProviderEntry) agent.ProviderConfig { // activeProviderIndex resolves the primary provider profile by ActiveProfile // id; list position is meaningless, so an unset or unknown id selects index 0. -func activeProviderIndex(option *Option) int { +func activeProviderIndex(option *cfg.Option) int { if option.ActiveProfile != "" { for i, entry := range option.Providers { if entry.ID == option.ActiveProfile { @@ -68,16 +49,16 @@ func activeProviderIndex(option *Option) int { return 0 } -func applyProviderLimits(cfg *agent.ProviderConfig, option *Option) { +func applyProviderLimits(providerConfig *agent.ProviderConfig, option *cfg.Option) { if option.MaxTokens != 0 { - cfg.MaxTokens = option.MaxTokens + providerConfig.MaxTokens = option.MaxTokens } if option.ContextWindow != 0 { - cfg.ContextWindow = option.ContextWindow + providerConfig.ContextWindow = option.ContextWindow } } -func ProviderConfig(option *Option) agent.ProviderConfig { +func ProviderConfig(option *cfg.Option) agent.ProviderConfig { if !hasSingleProviderFields(option) && len(option.Providers) > 0 { cfg := entryToProviderConfig(option.Providers[activeProviderIndex(option)]) applyProviderLimits(&cfg, option) @@ -107,7 +88,7 @@ func ProviderConfig(option *Option) agent.ProviderConfig { return cfg } -func FallbackProviderConfigs(option *Option) []agent.ProviderConfig { +func FallbackProviderConfigs(option *cfg.Option) []agent.ProviderConfig { if !hasSingleProviderFields(option) && len(option.Providers) > 0 { active := activeProviderIndex(option) var configs []agent.ProviderConfig @@ -126,11 +107,11 @@ func FallbackProviderConfigs(option *Option) []agent.ProviderConfig { return configs } -func ApplyResolvedProviderOptions(option *Option, cfg agent.ProviderConfig) { - option.Provider = cfg.Provider - option.BaseURL = cfg.BaseURL - option.APIKey = cfg.APIKey - option.Model = cfg.Model - option.MaxTokens = cfg.MaxTokens - option.ContextWindow = cfg.ContextWindow +func ApplyResolvedProviderOptions(option *cfg.Option, providerConfig agent.ProviderConfig) { + option.Provider = providerConfig.Provider + option.BaseURL = providerConfig.BaseURL + option.APIKey = providerConfig.APIKey + option.Model = providerConfig.Model + option.MaxTokens = providerConfig.MaxTokens + option.ContextWindow = providerConfig.ContextWindow } diff --git a/pkg/runner/provider_config_test.go b/pkg/runner/provider_config_test.go new file mode 100644 index 00000000..9476676b --- /dev/null +++ b/pkg/runner/provider_config_test.go @@ -0,0 +1,39 @@ +package runner + +import ( + "testing" + + cfg "github.com/chainreactors/aiscan/core/config" +) + +func TestProviderConfigSelectsActiveProfileAndFallbacks(t *testing.T) { + option := cfg.Option{LLMOptions: cfg.LLMOptions{ + ActiveProfile: "openai", + Providers: []cfg.LLMProviderEntry{ + {ID: "deepseek", Provider: "deepseek", APIKey: "dk-111", Model: "deepseek-chat", MaxTokens: 8192}, + {ID: "openai", Provider: "openai", APIKey: "sk-222", Model: "gpt-4o", MaxTokens: 32768}, + }, + }} + primary := ProviderConfig(&option) + if primary.Provider != "openai" || primary.APIKey != "sk-222" || primary.MaxTokens != 32768 { + t.Fatalf("primary profile = %+v", primary) + } + fallbacks := FallbackProviderConfigs(&option) + if len(fallbacks) != 1 || fallbacks[0].Provider != "deepseek" || fallbacks[0].APIKey != "dk-111" { + t.Fatalf("fallback profiles = %+v", fallbacks) + } +} + +func TestProviderConfigExplicitFieldsWin(t *testing.T) { + option := cfg.Option{LLMOptions: cfg.LLMOptions{ + Provider: "anthropic", APIKey: "cli-key", Model: "cli-model", + Providers: []cfg.LLMProviderEntry{{Provider: "deepseek", APIKey: "fallback-key", Model: "deepseek-chat"}}, + }} + primary := ProviderConfig(&option) + if primary.Provider != "anthropic" || primary.APIKey != "cli-key" || primary.Model != "cli-model" { + t.Fatalf("explicit provider = %+v", primary) + } + if fallbacks := FallbackProviderConfigs(&option); len(fallbacks) != 1 || fallbacks[0].Provider != "deepseek" { + t.Fatalf("fallback profiles = %+v", fallbacks) + } +} diff --git a/core/runner/remote_repl.go b/pkg/runner/remote_repl.go similarity index 100% rename from core/runner/remote_repl.go rename to pkg/runner/remote_repl.go diff --git a/core/runner/remote_repl_test.go b/pkg/runner/remote_repl_test.go similarity index 99% rename from core/runner/remote_repl_test.go rename to pkg/runner/remote_repl_test.go index 6e49bb4e..7a8ca656 100644 --- a/core/runner/remote_repl_test.go +++ b/pkg/runner/remote_repl_test.go @@ -7,7 +7,7 @@ import ( "time" cfg "github.com/chainreactors/aiscan/core/config" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/utils/pty" ) diff --git a/core/runner/runner.go b/pkg/runner/runner.go similarity index 96% rename from core/runner/runner.go rename to pkg/runner/runner.go index a0df525b..20d15b00 100644 --- a/core/runner/runner.go +++ b/pkg/runner/runner.go @@ -9,15 +9,15 @@ import ( "sync" "time" + "github.com/chainreactors/aiscan/agent" + inboxpkg "github.com/chainreactors/aiscan/agent/inbox" + tmuxpkg "github.com/chainreactors/aiscan/agent/tmux" + "github.com/chainreactors/aiscan/core/aop" cfg "github.com/chainreactors/aiscan/core/config" "github.com/chainreactors/aiscan/core/eventbus" + "github.com/chainreactors/aiscan/core/telemetry" coretool "github.com/chainreactors/aiscan/core/tool" - "github.com/chainreactors/aiscan/pkg/agent" - inboxpkg "github.com/chainreactors/aiscan/pkg/agent/inbox" - tmuxpkg "github.com/chainreactors/aiscan/pkg/agent/tmux" - "github.com/chainreactors/aiscan/pkg/aop" cmdpkg "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" "github.com/chainreactors/aiscan/pkg/tui" "github.com/chainreactors/aiscan/skills" "github.com/chainreactors/aiscan/tools/toolargs" @@ -67,7 +67,7 @@ const ( type RuntimeConfig struct { ExistingApp *App - IOA *cfg.IOAConfig + IOA *IOAConfig PromptConfig *PromptConfig NoOutput bool InteractiveOutput bool @@ -107,7 +107,7 @@ func NewAgentRuntime(ctx context.Context, option *cfg.Option, logger telemetry.L rt.app = rc.ExistingApp } else { providerOptional := rc != nil && (rc.IOA != nil || rc.ProviderOptional) - appCfg := cfg.AppConfig(option, cfg.RuntimeFeatures{ + appCfg := AppConfig(option, RuntimeFeatures{ ProviderEnabled: true, ProviderOptional: providerOptional, ToolsEnabled: true, @@ -122,7 +122,7 @@ func NewAgentRuntime(ctx context.Context, option *cfg.Option, logger telemetry.L } rt.app = application rt.ownsApp = true - cfg.ApplyResolvedProviderOptions(option, application.ProviderConfig) + ApplyResolvedProviderOptions(option, application.ProviderConfig) for _, d := range application.SkillDiagnostics { logger.Warnf("skill %s: %s", d.Path, d.Message) @@ -227,6 +227,7 @@ func NewAgentRuntime(ctx context.Context, option *cfg.Option, logger telemetry.L Logger: logger, CacheRetention: agent.CacheShort, Bus: rt.kernelBus, + Hooks: rt.app.Hooks, } if option.SaveSession { @@ -269,7 +270,7 @@ func NewAgentRuntime(ctx context.Context, option *cfg.Option, logger telemetry.L } handoffCancel = subscribeIOAHandoffContext(rt.ctx, publicBus, rt.app.IOAClient, ioaSpace, logger) rt.app.Commands.RegisterTool(subAgentTool) - loop := agent.NewLoopCommand() + loop := newLoopCommand() rt.app.Commands.Register(cmdpkg.Command{Name: loop.Name(), Usage: loop.Usage(), Run: loop.Run}, "loop") if option.Resume != "" { @@ -384,7 +385,7 @@ func (rt *AgentRuntime) ReloadProvider(option *cfg.Option) (agent.Provider, stri if logger == nil { logger = telemetry.NopLogger() } - provider, resolved, err := initProvider(cfg.ProviderConfig(option), logger) + provider, resolved, err := initProvider(ProviderConfig(option), logger) if err != nil { return nil, "", err } @@ -543,7 +544,7 @@ func RunDirectScannerMode(ctx context.Context, option *cfg.Option, rest []string defer restoreLogs() } - application, err := NewApp(ctx, cfg.AppConfig(option, features, scannerLogger)) + application, err := NewApp(ctx, AppConfig(option, features, scannerLogger)) if err != nil { return fmt.Errorf("init app: %w", err) } @@ -551,7 +552,7 @@ func RunDirectScannerMode(ctx context.Context, option *cfg.Option, rest []string if err := application.WaitEngines(ctx); err != nil { return fmt.Errorf("engine init: %w", err) } - cfg.ApplyResolvedProviderOptions(option, application.ProviderConfig) + ApplyResolvedProviderOptions(option, application.ProviderConfig) if !application.Commands.Has(scannerArgs[0]) { return fmt.Errorf("unknown subcommand: %s", scannerArgs[0]) @@ -691,7 +692,7 @@ func registerIOATools(ctx context.Context, application *App, option *cfg.Option) if ioaURL == "" { return nil } - ioaCfg := cfg.IOAConfig{ + ioaCfg := IOAConfig{ URL: ioaURL, NodeID: option.IOANodeID, NodeName: option.IOANodeName, diff --git a/pkg/runner/runtime_config.go b/pkg/runner/runtime_config.go new file mode 100644 index 00000000..3ef69b83 --- /dev/null +++ b/pkg/runner/runtime_config.go @@ -0,0 +1,18 @@ +package runner + +import ( + "github.com/chainreactors/aiscan/agent" + cfg "github.com/chainreactors/aiscan/core/config" +) + +// ResolveRuntimeConfig resolves the process configuration and applies process +// state such as the data directory. +func ResolveRuntimeConfig(option *cfg.Option) (string, error) { + return cfg.ResolveRuntimeConfig(option, true, agent.InferProviderFromBaseURL) +} + +// ResolveRuntimeConfigCandidate resolves a staged Web configuration without +// mutating process-wide state before the candidate is committed. +func ResolveRuntimeConfigCandidate(option *cfg.Option) (string, error) { + return cfg.ResolveRuntimeConfig(option, false, agent.InferProviderFromBaseURL) +} diff --git a/core/runner/runtime_protocol.go b/pkg/runner/runtime_protocol.go similarity index 100% rename from core/runner/runtime_protocol.go rename to pkg/runner/runtime_protocol.go diff --git a/core/runner/runtime_protocol_test.go b/pkg/runner/runtime_protocol_test.go similarity index 100% rename from core/runner/runtime_protocol_test.go rename to pkg/runner/runtime_protocol_test.go diff --git a/core/runner/runtime_semantics_test.go b/pkg/runner/runtime_semantics_test.go similarity index 93% rename from core/runner/runtime_semantics_test.go rename to pkg/runner/runtime_semantics_test.go index e2b1c4e3..7d5a7166 100644 --- a/core/runner/runtime_semantics_test.go +++ b/pkg/runner/runtime_semantics_test.go @@ -8,12 +8,13 @@ import ( "testing" "time" - "github.com/chainreactors/aiscan/pkg/agent" - "github.com/chainreactors/aiscan/pkg/agent/inbox" - "github.com/chainreactors/aiscan/pkg/aop" - xcommand "github.com/chainreactors/aiscan/pkg/aop/x/command" + "github.com/chainreactors/aiscan/agent" + "github.com/chainreactors/aiscan/agent/inbox" + "github.com/chainreactors/aiscan/core/aop" + xcommand "github.com/chainreactors/aiscan/core/aop/x/command" + "github.com/chainreactors/aiscan/core/capability" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" ) type runtimeSemanticProvider struct { @@ -155,7 +156,7 @@ func TestSessionContextCancellationStopsActiveRun(t *testing.T) { func TestCommandAddsAOPHistoryWithoutChangingTranscript(t *testing.T) { registry := commands.NewRegistry() - commands.BuildGroup("core", &commands.Deps{WorkDir: t.TempDir(), BashTimeout: 5, Logger: telemetry.NopLogger()}, registry) + commands.BuildPlan(capability.Select(capability.Options{Groups: []string{"core"}}), &commands.Deps{WorkDir: t.TempDir(), BashTimeout: 5, Logger: telemetry.NopLogger()}, registry) rt := newBareRuntime(t, registry, nil) session, err := rt.OpenSession(context.Background(), SessionOptions{ID: "session-1"}) if err != nil { diff --git a/core/runner/runtime_session.go b/pkg/runner/runtime_session.go similarity index 98% rename from core/runner/runtime_session.go rename to pkg/runner/runtime_session.go index 747f17ec..f114c163 100644 --- a/core/runner/runtime_session.go +++ b/pkg/runner/runtime_session.go @@ -9,14 +9,14 @@ import ( "sync" "time" + "github.com/chainreactors/aiscan/agent" + "github.com/chainreactors/aiscan/agent/evaluator" + inboxpkg "github.com/chainreactors/aiscan/agent/inbox" + "github.com/chainreactors/aiscan/core/aop" + xcommand "github.com/chainreactors/aiscan/core/aop/x/command" "github.com/chainreactors/aiscan/core/eventbus" - "github.com/chainreactors/aiscan/pkg/agent" - "github.com/chainreactors/aiscan/pkg/agent/evaluator" - inboxpkg "github.com/chainreactors/aiscan/pkg/agent/inbox" - "github.com/chainreactors/aiscan/pkg/aop" - xcommand "github.com/chainreactors/aiscan/pkg/aop/x/command" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" "github.com/chainreactors/aiscan/pkg/tui" "github.com/chainreactors/aiscan/skills" ) @@ -184,7 +184,7 @@ func (s *commandSession) execute(ctx context.Context, input string) commandOutco if line == "/continue" || strings.HasPrefix(line, "/followup ") || strings.HasPrefix(line, "/skill:") { return commandOutcome{err: fmt.Errorf("%s requires a Run", line)} } - ctx = commands.ContextWithInbox(ctx, s.state.inbox) + ctx = inboxpkg.ContextWithInbox(ctx, s.state.inbox) ctx = agent.ContextWithLoopScheduler(ctx, s.state.scheduler) if strings.HasPrefix(line, "!") { diff --git a/core/runner/runtime_session_isolation_test.go b/pkg/runner/runtime_session_isolation_test.go similarity index 87% rename from core/runner/runtime_session_isolation_test.go rename to pkg/runner/runtime_session_isolation_test.go index 14b328d1..ec67b017 100644 --- a/core/runner/runtime_session_isolation_test.go +++ b/pkg/runner/runtime_session_isolation_test.go @@ -6,12 +6,13 @@ import ( "testing" "time" + "github.com/chainreactors/aiscan/agent" + "github.com/chainreactors/aiscan/core/aop" + "github.com/chainreactors/aiscan/core/capability" cfg "github.com/chainreactors/aiscan/core/config" "github.com/chainreactors/aiscan/core/eventbus" - "github.com/chainreactors/aiscan/pkg/agent" - "github.com/chainreactors/aiscan/pkg/aop" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" ) func newBareRuntime(t *testing.T, reg *commands.CommandRegistry, provider agent.Provider) *AgentRuntime { @@ -36,8 +37,8 @@ func newBareRuntime(t *testing.T, reg *commands.CommandRegistry, provider agent. func TestRuntimeSessionDirectLoopUsesSessionScheduler(t *testing.T) { reg := commands.NewRegistry() - commands.BuildGroup("core", &commands.Deps{WorkDir: t.TempDir(), BashTimeout: 5, Logger: telemetry.NopLogger()}, reg) - loop := agent.NewLoopCommand() + commands.BuildPlan(capability.Select(capability.Options{Groups: []string{"core"}}), &commands.Deps{WorkDir: t.TempDir(), BashTimeout: 5, Logger: telemetry.NopLogger()}, reg) + loop := newLoopCommand() reg.Register(commands.Command{Name: loop.Name(), Usage: loop.Usage(), Run: loop.Run}, "loop") rt := newBareRuntime(t, reg, nil) t.Cleanup(func() { diff --git a/core/runner/scanner.go b/pkg/runner/scanner.go similarity index 90% rename from core/runner/scanner.go rename to pkg/runner/scanner.go index 5608e181..69862489 100644 --- a/core/runner/scanner.go +++ b/pkg/runner/scanner.go @@ -7,19 +7,19 @@ import ( "github.com/chainreactors/aiscan/core/config" ) -func DirectScannerRuntimeFeatures(rest []string) (config.RuntimeFeatures, []string, error) { +func DirectScannerRuntimeFeatures(rest []string) (RuntimeFeatures, []string, error) { if len(rest) == 0 { - return config.RuntimeFeatures{}, nil, fmt.Errorf("missing scanner command") + return RuntimeFeatures{}, nil, fmt.Errorf("missing scanner command") } if rest[0] != "scan" { - return config.RuntimeFeatures{}, rest, nil + return RuntimeFeatures{}, rest, nil } verifyMode, explicit := scannerVerifyMode(rest[1:]) sniperEnabled := HasScannerFlag(rest[1:], "--sniper") deepEnabled := HasScannerFlag(rest[1:], "--deep") aiSkillRequested := sniperEnabled || deepEnabled - features := config.RuntimeFeatures{} + features := RuntimeFeatures{} if aiSkillRequested { features.ProviderEnabled = true @@ -52,7 +52,7 @@ func DirectScannerRuntimeFeatures(rest []string) (config.RuntimeFeatures, []stri return features, rest, nil default: if explicit { - return config.RuntimeFeatures{}, nil, fmt.Errorf("invalid --verify value %q: expected auto, off, low, medium, high, or critical", verifyMode) + return RuntimeFeatures{}, nil, fmt.Errorf("invalid --verify value %q: expected auto, off, low, medium, high, or critical", verifyMode) } return features, rest, nil } diff --git a/core/runner/stdio.go b/pkg/runner/stdio.go similarity index 96% rename from core/runner/stdio.go rename to pkg/runner/stdio.go index aa7c4516..c7c9d34a 100644 --- a/core/runner/stdio.go +++ b/pkg/runner/stdio.go @@ -9,9 +9,9 @@ import ( "strings" "sync" + "github.com/chainreactors/aiscan/core/aop" cfg "github.com/chainreactors/aiscan/core/config" - "github.com/chainreactors/aiscan/pkg/aop" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/webproto" ) diff --git a/core/runner/stdio_concurrency_test.go b/pkg/runner/stdio_concurrency_test.go similarity index 98% rename from core/runner/stdio_concurrency_test.go rename to pkg/runner/stdio_concurrency_test.go index d8729d6c..42100898 100644 --- a/core/runner/stdio_concurrency_test.go +++ b/pkg/runner/stdio_concurrency_test.go @@ -8,8 +8,8 @@ import ( "testing" "time" - "github.com/chainreactors/aiscan/pkg/agent" - "github.com/chainreactors/aiscan/pkg/aop" + "github.com/chainreactors/aiscan/agent" + "github.com/chainreactors/aiscan/core/aop" "github.com/chainreactors/aiscan/pkg/webproto" ) diff --git a/core/runner/stdio_test.go b/pkg/runner/stdio_test.go similarity index 98% rename from core/runner/stdio_test.go rename to pkg/runner/stdio_test.go index a09b1b67..d13edce0 100644 --- a/core/runner/stdio_test.go +++ b/pkg/runner/stdio_test.go @@ -9,8 +9,8 @@ import ( "strings" "testing" - "github.com/chainreactors/aiscan/pkg/aop" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/core/aop" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/webproto" ) diff --git a/core/runner/subagent_handoff.go b/pkg/runner/subagent_handoff.go similarity index 97% rename from core/runner/subagent_handoff.go rename to pkg/runner/subagent_handoff.go index 6650676d..79b60442 100644 --- a/core/runner/subagent_handoff.go +++ b/pkg/runner/subagent_handoff.go @@ -8,11 +8,11 @@ import ( "sync" "time" + "github.com/chainreactors/aiscan/agent" + "github.com/chainreactors/aiscan/core/aop" + "github.com/chainreactors/aiscan/core/aop/x/delegation" "github.com/chainreactors/aiscan/core/eventbus" - "github.com/chainreactors/aiscan/pkg/agent" - "github.com/chainreactors/aiscan/pkg/aop" - "github.com/chainreactors/aiscan/pkg/aop/x/delegation" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/ioa/protocols" ) diff --git a/core/runner/subagent_handoff_test.go b/pkg/runner/subagent_handoff_test.go similarity index 98% rename from core/runner/subagent_handoff_test.go rename to pkg/runner/subagent_handoff_test.go index 77adc5e5..554b4598 100644 --- a/core/runner/subagent_handoff_test.go +++ b/pkg/runner/subagent_handoff_test.go @@ -7,9 +7,9 @@ import ( "testing" "time" + "github.com/chainreactors/aiscan/core/aop" + "github.com/chainreactors/aiscan/core/aop/x/delegation" "github.com/chainreactors/aiscan/core/eventbus" - "github.com/chainreactors/aiscan/pkg/aop" - "github.com/chainreactors/aiscan/pkg/aop/x/delegation" "github.com/chainreactors/ioa/protocols" ) diff --git a/core/transport/transport.go b/pkg/transport/transport.go similarity index 89% rename from core/transport/transport.go rename to pkg/transport/transport.go index a5cac729..cede239f 100644 --- a/core/transport/transport.go +++ b/pkg/transport/transport.go @@ -5,8 +5,8 @@ import ( "io" cfg "github.com/chainreactors/aiscan/core/config" - "github.com/chainreactors/aiscan/core/runner" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/core/telemetry" + "github.com/chainreactors/aiscan/pkg/runner" "github.com/chainreactors/aiscan/pkg/webagent" ) diff --git a/pkg/tui/banner_render_test.go b/pkg/tui/banner_render_test.go index 11ed20f5..486a4188 100644 --- a/pkg/tui/banner_render_test.go +++ b/pkg/tui/banner_render_test.go @@ -4,8 +4,8 @@ import ( "strings" "testing" + "github.com/chainreactors/aiscan/agent" outputpkg "github.com/chainreactors/aiscan/core/output" - "github.com/chainreactors/aiscan/pkg/agent" ) // assertUniformWidth checks every line of a rendered box has the same visible diff --git a/pkg/tui/commands.go b/pkg/tui/commands.go index 925b72b1..75dee449 100644 --- a/pkg/tui/commands.go +++ b/pkg/tui/commands.go @@ -6,10 +6,10 @@ import ( "net/url" "strings" + "github.com/chainreactors/aiscan/agent" cfg "github.com/chainreactors/aiscan/core/config" - "github.com/chainreactors/aiscan/pkg/agent" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" "github.com/chainreactors/aiscan/pkg/webproto" "github.com/chainreactors/aiscan/skills" ) diff --git a/pkg/tui/console.go b/pkg/tui/console.go index 0dc558e5..85c764bd 100644 --- a/pkg/tui/console.go +++ b/pkg/tui/console.go @@ -16,12 +16,12 @@ import ( "time" "github.com/carapace-sh/carapace" + "github.com/chainreactors/aiscan/agent" + "github.com/chainreactors/aiscan/agent/probe" cfg "github.com/chainreactors/aiscan/core/config" outputpkg "github.com/chainreactors/aiscan/core/output" - "github.com/chainreactors/aiscan/pkg/agent" - "github.com/chainreactors/aiscan/pkg/agent/probe" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" ioaclient "github.com/chainreactors/ioa/client" "github.com/chainreactors/tui/console" rlterm "github.com/chainreactors/tui/readline/terminal" @@ -1466,7 +1466,12 @@ func (r *AgentConsole) applyProviderConfig(pc agent.ProviderConfig) (agent.Provi } r.output.SetContextWindow(contextWindow) if r.option != nil { - cfg.ApplyResolvedProviderOptions(r.option, *resolved) + r.option.Provider = resolved.Provider + r.option.BaseURL = resolved.BaseURL + r.option.APIKey = resolved.APIKey + r.option.Model = resolved.Model + r.option.MaxTokens = resolved.MaxTokens + r.option.ContextWindow = resolved.ContextWindow r.option.LLMProxy = resolved.Proxy } r.syncEvalToController() diff --git a/pkg/tui/console_test.go b/pkg/tui/console_test.go index 1689c6c4..52d8cf57 100644 --- a/pkg/tui/console_test.go +++ b/pkg/tui/console_test.go @@ -14,8 +14,8 @@ import ( "testing" "time" + "github.com/chainreactors/aiscan/agent" cfg "github.com/chainreactors/aiscan/core/config" - "github.com/chainreactors/aiscan/pkg/agent" "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/tui/readline/inputrc" rlterm "github.com/chainreactors/tui/readline/terminal" diff --git a/pkg/tui/controller.go b/pkg/tui/controller.go index c6862736..d0d1b998 100644 --- a/pkg/tui/controller.go +++ b/pkg/tui/controller.go @@ -7,9 +7,9 @@ import ( "strings" "sync" - "github.com/chainreactors/aiscan/pkg/agent" - "github.com/chainreactors/aiscan/pkg/agent/evaluator" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/agent" + "github.com/chainreactors/aiscan/agent/evaluator" + "github.com/chainreactors/aiscan/core/telemetry" ) type agentRunFunc func(context.Context) (*agent.Result, error) diff --git a/pkg/tui/controller_test.go b/pkg/tui/controller_test.go index 92cef548..8717d7b9 100644 --- a/pkg/tui/controller_test.go +++ b/pkg/tui/controller_test.go @@ -8,7 +8,7 @@ import ( "testing" "time" - "github.com/chainreactors/aiscan/pkg/agent" + "github.com/chainreactors/aiscan/agent" ) // gateProvider blocks the first call until released; later calls answer diff --git a/pkg/tui/format.go b/pkg/tui/format.go index 5e66cdd3..a60a0f4b 100644 --- a/pkg/tui/format.go +++ b/pkg/tui/format.go @@ -11,10 +11,10 @@ import ( "unicode" "unicode/utf8" - "github.com/chainreactors/aiscan/pkg/agent" - "github.com/chainreactors/aiscan/pkg/agent/truncate" - "github.com/chainreactors/aiscan/pkg/aop" - "github.com/chainreactors/aiscan/pkg/util" + "github.com/chainreactors/aiscan/agent" + "github.com/chainreactors/aiscan/core/aop" + "github.com/chainreactors/aiscan/core/truncate" + "github.com/chainreactors/aiscan/core/util" "github.com/charmbracelet/glamour" "github.com/muesli/termenv" "golang.org/x/term" diff --git a/pkg/tui/live.go b/pkg/tui/live.go index 20735502..b981b08e 100644 --- a/pkg/tui/live.go +++ b/pkg/tui/live.go @@ -5,9 +5,9 @@ import ( "strings" "time" - "github.com/chainreactors/aiscan/pkg/agent" - "github.com/chainreactors/aiscan/pkg/agent/truncate" - "github.com/chainreactors/aiscan/pkg/util" + "github.com/chainreactors/aiscan/agent" + "github.com/chainreactors/aiscan/core/truncate" + "github.com/chainreactors/aiscan/core/util" ) const ( diff --git a/pkg/tui/output.go b/pkg/tui/output.go index 654a826a..f605d79a 100644 --- a/pkg/tui/output.go +++ b/pkg/tui/output.go @@ -9,14 +9,14 @@ import ( "sync" "time" + "github.com/chainreactors/aiscan/agent" + "github.com/chainreactors/aiscan/core/aop" + xcompact "github.com/chainreactors/aiscan/core/aop/x/compact" + xeval "github.com/chainreactors/aiscan/core/aop/x/eval" cfg "github.com/chainreactors/aiscan/core/config" "github.com/chainreactors/aiscan/core/output" - "github.com/chainreactors/aiscan/pkg/agent" - "github.com/chainreactors/aiscan/pkg/agent/truncate" - "github.com/chainreactors/aiscan/pkg/aop" - xcompact "github.com/chainreactors/aiscan/pkg/aop/x/compact" - xeval "github.com/chainreactors/aiscan/pkg/aop/x/eval" - "github.com/chainreactors/aiscan/pkg/util" + "github.com/chainreactors/aiscan/core/truncate" + "github.com/chainreactors/aiscan/core/util" "golang.org/x/term" ) diff --git a/pkg/tui/output_test.go b/pkg/tui/output_test.go index 5889e3f8..3561c038 100644 --- a/pkg/tui/output_test.go +++ b/pkg/tui/output_test.go @@ -10,11 +10,11 @@ import ( "testing" "time" + "github.com/chainreactors/aiscan/agent" + "github.com/chainreactors/aiscan/core/aop" + xeval "github.com/chainreactors/aiscan/core/aop/x/eval" cfg "github.com/chainreactors/aiscan/core/config" "github.com/chainreactors/aiscan/core/output" - "github.com/chainreactors/aiscan/pkg/agent" - "github.com/chainreactors/aiscan/pkg/aop" - xeval "github.com/chainreactors/aiscan/pkg/aop/x/eval" ) type syncedBuffer struct { diff --git a/pkg/tui/remote_console.go b/pkg/tui/remote_console.go index dd8193f1..1a03fcbd 100644 --- a/pkg/tui/remote_console.go +++ b/pkg/tui/remote_console.go @@ -7,9 +7,9 @@ import ( "io" "sync" + "github.com/chainreactors/aiscan/agent" + "github.com/chainreactors/aiscan/core/aop" cfg "github.com/chainreactors/aiscan/core/config" - "github.com/chainreactors/aiscan/pkg/agent" - "github.com/chainreactors/aiscan/pkg/aop" rlterm "github.com/chainreactors/tui/readline/terminal" ) diff --git a/pkg/tui/remote_console_test.go b/pkg/tui/remote_console_test.go index 6822e751..d4433705 100644 --- a/pkg/tui/remote_console_test.go +++ b/pkg/tui/remote_console_test.go @@ -4,8 +4,8 @@ import ( "bytes" "testing" - "github.com/chainreactors/aiscan/pkg/agent" - "github.com/chainreactors/aiscan/pkg/aop" + "github.com/chainreactors/aiscan/agent" + "github.com/chainreactors/aiscan/core/aop" ) func TestSubscribeAgentOutputRestoresSessionEvents(t *testing.T) { diff --git a/pkg/tui/render.go b/pkg/tui/render.go index b7c1ed40..46d19dcb 100644 --- a/pkg/tui/render.go +++ b/pkg/tui/render.go @@ -8,7 +8,7 @@ import ( "sync" "time" - "github.com/chainreactors/aiscan/pkg/util" + "github.com/chainreactors/aiscan/core/util" bspinner "github.com/charmbracelet/bubbles/spinner" ) diff --git a/pkg/web/agents.go b/pkg/web/agents.go index 38ac936d..95d14491 100644 --- a/pkg/web/agents.go +++ b/pkg/web/agents.go @@ -9,8 +9,8 @@ import ( "sync/atomic" "time" + "github.com/chainreactors/aiscan/core/aop" "github.com/chainreactors/aiscan/core/output" - "github.com/chainreactors/aiscan/pkg/aop" "github.com/chainreactors/aiscan/pkg/webproto" "github.com/chainreactors/ioa/protocols" "github.com/chainreactors/utils/pty" diff --git a/pkg/web/agents_session_end_test.go b/pkg/web/agents_session_end_test.go index 30d9e6ea..e4ae3887 100644 --- a/pkg/web/agents_session_end_test.go +++ b/pkg/web/agents_session_end_test.go @@ -5,7 +5,7 @@ import ( "testing" "time" - "github.com/chainreactors/aiscan/pkg/aop" + "github.com/chainreactors/aiscan/core/aop" ) func sessionEvent(t *testing.T, typ, sessionID string, data any) aop.Event { diff --git a/pkg/web/agents_test.go b/pkg/web/agents_test.go index 4d81151e..5beb8aed 100644 --- a/pkg/web/agents_test.go +++ b/pkg/web/agents_test.go @@ -15,8 +15,8 @@ import ( webstatic "github.com/chainreactors/aiscan/web" + "github.com/chainreactors/aiscan/core/aop" "github.com/chainreactors/aiscan/core/output" - "github.com/chainreactors/aiscan/pkg/aop" "github.com/chainreactors/aiscan/pkg/webproto" "github.com/chainreactors/ioa/protocols" "github.com/chainreactors/utils/pty" diff --git a/pkg/web/config_transaction_test.go b/pkg/web/config_transaction_test.go index 9a5dd9e3..c2f3a9c4 100644 --- a/pkg/web/config_transaction_test.go +++ b/pkg/web/config_transaction_test.go @@ -7,7 +7,7 @@ import ( "testing" "time" - "github.com/chainreactors/aiscan/core/runner" + "github.com/chainreactors/aiscan/pkg/runner" "github.com/chainreactors/aiscan/pkg/webproto" ) diff --git a/pkg/web/conn_probe_test.go b/pkg/web/conn_probe_test.go index f7c91d5d..7b956aa0 100644 --- a/pkg/web/conn_probe_test.go +++ b/pkg/web/conn_probe_test.go @@ -8,7 +8,7 @@ import ( "strings" "testing" - "github.com/chainreactors/aiscan/pkg/agent/probe" + "github.com/chainreactors/aiscan/pkg/probe" "github.com/chainreactors/aiscan/pkg/webproto" ) diff --git a/pkg/web/eval_forward_test.go b/pkg/web/eval_forward_test.go index 59f97f74..033cbdb5 100644 --- a/pkg/web/eval_forward_test.go +++ b/pkg/web/eval_forward_test.go @@ -4,9 +4,9 @@ import ( "encoding/json" "testing" - "github.com/chainreactors/aiscan/pkg/aop" - xcompact "github.com/chainreactors/aiscan/pkg/aop/x/compact" - xeval "github.com/chainreactors/aiscan/pkg/aop/x/eval" + "github.com/chainreactors/aiscan/core/aop" + xcompact "github.com/chainreactors/aiscan/core/aop/x/compact" + xeval "github.com/chainreactors/aiscan/core/aop/x/eval" "github.com/chainreactors/aiscan/pkg/webproto" ) diff --git a/pkg/web/handler.go b/pkg/web/handler.go index 344a692d..602342e3 100644 --- a/pkg/web/handler.go +++ b/pkg/web/handler.go @@ -9,7 +9,7 @@ import ( "strconv" "strings" - "github.com/chainreactors/aiscan/pkg/agent/probe" + "github.com/chainreactors/aiscan/agent/probe" "github.com/chainreactors/aiscan/pkg/webproto" ) diff --git a/pkg/web/llm_probe_test.go b/pkg/web/llm_probe_test.go index 5f1c012d..ead02193 100644 --- a/pkg/web/llm_probe_test.go +++ b/pkg/web/llm_probe_test.go @@ -8,7 +8,7 @@ import ( "strings" "testing" - "github.com/chainreactors/aiscan/pkg/agent/probe" + "github.com/chainreactors/aiscan/agent/probe" "github.com/chainreactors/aiscan/pkg/webproto" ) diff --git a/pkg/web/probe.go b/pkg/web/probe.go index e8712689..d125c7b2 100644 --- a/pkg/web/probe.go +++ b/pkg/web/probe.go @@ -4,7 +4,8 @@ import ( "context" "strings" - "github.com/chainreactors/aiscan/pkg/agent/probe" + agentprobe "github.com/chainreactors/aiscan/agent/probe" + "github.com/chainreactors/aiscan/pkg/probe" "github.com/chainreactors/aiscan/pkg/webproto" ) @@ -30,15 +31,15 @@ func toProbeConfig(dc webproto.DistributeConfig) probe.ProbeConfig { // TestLLM probes the supplied LLM settings, falling back to the stored API key // when the request leaves it blank, then delegates to pkg/probe. -func (s *Service) TestLLM(ctx context.Context, req probe.LLMProbeRequest) (probe.LLMTestResult, error) { - return probe.TestLLM(ctx, req, s.storedLLMAPIKey(ctx, req.ProfileID)) +func (s *Service) TestLLM(ctx context.Context, req agentprobe.LLMProbeRequest) (agentprobe.LLMTestResult, error) { + return agentprobe.TestLLM(ctx, req, s.storedLLMAPIKey(ctx, req.ProfileID)) } // ListLLMModels enumerates the models the supplied LLM endpoint advertises, // falling back to the stored API key when the request leaves it blank, then // delegates to pkg/probe. -func (s *Service) ListLLMModels(ctx context.Context, req probe.LLMProbeRequest) (probe.LLMModelsResult, error) { - return probe.ListLLMModels(ctx, req, s.storedLLMAPIKey(ctx, req.ProfileID)) +func (s *Service) ListLLMModels(ctx context.Context, req agentprobe.LLMProbeRequest) (agentprobe.LLMModelsResult, error) { + return agentprobe.ListLLMModels(ctx, req, s.storedLLMAPIKey(ctx, req.ProfileID)) } // storedLLMAPIKey returns the requested profile's persisted API key. A blank diff --git a/pkg/web/replay_test.go b/pkg/web/replay_test.go index 6ea81db5..cbed9506 100644 --- a/pkg/web/replay_test.go +++ b/pkg/web/replay_test.go @@ -11,7 +11,7 @@ import ( "testing" "time" - "github.com/chainreactors/aiscan/pkg/aop" + "github.com/chainreactors/aiscan/core/aop" ) type lockedResponseRecorder struct { diff --git a/pkg/web/report.go b/pkg/web/report.go new file mode 100644 index 00000000..b2b8d74c --- /dev/null +++ b/pkg/web/report.go @@ -0,0 +1,18 @@ +package web + +import "github.com/chainreactors/aiscan/core/output" + +// defaultReportLang is the language the report is frozen in at scan time; the +// stored copy is only a fallback because GetReport re-renders per request. +const defaultReportLang = "zh" + +func buildMarkdownReport(target, mode string, result *output.Result, lang string) string { + return output.RenderReport(result, output.ReportOptions{ + Style: output.StyleMarkdown, + Lang: lang, + Title: target, + Mode: mode, + Sitemap: true, + CollapseBare: true, + }) +} diff --git a/pkg/web/scan_lifecycle_test.go b/pkg/web/scan_lifecycle_test.go index 6e00db86..4c00ca3c 100644 --- a/pkg/web/scan_lifecycle_test.go +++ b/pkg/web/scan_lifecycle_test.go @@ -137,6 +137,28 @@ func TestCancelQueuedScanDoesNotWaitForConcurrencySlot(t *testing.T) { waitScanStatus(t, store, running.ID, StatusCanceled) } +type controlledDeadlineContext struct { + context.Context + done chan struct{} +} + +func newControlledDeadlineContext() *controlledDeadlineContext { + return &controlledDeadlineContext{Context: context.Background(), done: make(chan struct{})} +} + +func (c *controlledDeadlineContext) Done() <-chan struct{} { return c.done } + +func (c *controlledDeadlineContext) Err() error { + select { + case <-c.done: + return context.DeadlineExceeded + default: + return nil + } +} + +func (c *controlledDeadlineContext) expire() { close(c.done) } + func TestRemoteScanTimeoutCancelsAgentAndFailsScan(t *testing.T) { store, err := NewSQLiteStore(filepath.Join(t.TempDir(), "web.db")) if err != nil { @@ -144,34 +166,94 @@ func TestRemoteScanTimeoutCancelsAgentAndFailsScan(t *testing.T) { } t.Cleanup(func() { _ = store.Close() }) - svc := NewService(ServiceConfig{Store: store, MaxConcurrent: 1, ScanTimeout: 50 * time.Millisecond}) + svc := NewService(ServiceConfig{Store: store, MaxConcurrent: 1}) pool := NewAgentPool(svc.Hub()) svc.SetAgentPool(pool) - srv, _ := setupTestServerWithPool(t, pool) - conn := dialAgent(t, srv, "timeout-agent", []string{"scan"}) - t.Cleanup(func() { _ = conn.Close() }) - waitAgents(t, pool, 1) + agent := newFakeAgent("timeout-agent", 1) + pool.register(agent) - job, err := svc.SubmitScan(context.Background(), "127.0.0.1", "quick", false, false, false) - if err != nil { + now := time.Now() + job := &ScanJob{ + ID: "timeout-scan", Target: "127.0.0.1", Mode: "quick", + Status: StatusRunning, CreatedAt: now, UpdatedAt: now, + } + if err := store.Create(context.Background(), job); err != nil { t.Fatal(err) } + jobID := job.ID + + ctx := newControlledDeadlineContext() + done := make(chan struct{}) + go func() { + svc.runScanViaAgent(ctx, job) + close(done) + }() + var call webproto.Message - if err := conn.ReadJSON(&call); err != nil { - t.Fatal(err) + select { + case call = <-agent.sendCh: + case <-time.After(time.Second): + t.Fatal("agent did not receive scan dispatch") } - _ = conn.SetReadDeadline(time.Now().Add(time.Second)) + if call.Type != webproto.TypeAOP || call.TaskID != jobID { + t.Fatalf("scan dispatch = %+v", call) + } + + ctx.expire() var cancel webproto.Message - if err := conn.ReadJSON(&cancel); err != nil { - t.Fatalf("agent did not receive timeout cancellation: %v", err) + select { + case cancel = <-agent.controlCh: + case <-time.After(time.Second): + t.Fatal("agent did not receive timeout cancellation") } - if cancel.Type != "cancel" || cancel.TaskID != job.ID { + if cancel.Type != "cancel" || cancel.TaskID != jobID { t.Fatalf("timeout cancel frame = %+v", cancel) } + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("timed-out remote scan did not return") + } + failed := waitScanStatus(t, store, jobID, StatusFailed) + if failed.Error != "scan timed out" { + t.Fatalf("timeout error = %q", failed.Error) + } +} + +func TestRemoteScanExpiredBeforeDispatchFailsScan(t *testing.T) { + store, err := NewSQLiteStore(filepath.Join(t.TempDir(), "web.db")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = store.Close() }) + + svc := NewService(ServiceConfig{Store: store, MaxConcurrent: 1}) + pool := NewAgentPool(svc.Hub()) + svc.SetAgentPool(pool) + agent := newFakeAgent("timeout-agent", 1) + pool.register(agent) + + now := time.Now() + job := &ScanJob{ + ID: "expired-scan", Target: "127.0.0.1", Mode: "quick", + Status: StatusRunning, CreatedAt: now, UpdatedAt: now, + } + if err := store.Create(context.Background(), job); err != nil { + t.Fatal(err) + } + ctx := newControlledDeadlineContext() + ctx.expire() + svc.runScanViaAgent(ctx, job) + failed := waitScanStatus(t, store, job.ID, StatusFailed) if failed.Error != "scan timed out" { t.Fatalf("timeout error = %q", failed.Error) } + select { + case msg := <-agent.sendCh: + t.Fatalf("expired scan was dispatched: %+v", msg) + default: + } } func setupTestServerWithPool(t *testing.T, pool *AgentPool) (*httptest.Server, *AgentPool) { diff --git a/pkg/web/service.go b/pkg/web/service.go index 9b49b04e..aee07d76 100644 --- a/pkg/web/service.go +++ b/pkg/web/service.go @@ -17,13 +17,13 @@ import ( "sync" "time" + "github.com/chainreactors/aiscan/core/aop" + xcompact "github.com/chainreactors/aiscan/core/aop/x/compact" + xeval "github.com/chainreactors/aiscan/core/aop/x/eval" "github.com/chainreactors/aiscan/core/config" "github.com/chainreactors/aiscan/core/output" - "github.com/chainreactors/aiscan/core/runner" - "github.com/chainreactors/aiscan/pkg/aop" - xcompact "github.com/chainreactors/aiscan/pkg/aop/x/compact" - xeval "github.com/chainreactors/aiscan/pkg/aop/x/eval" "github.com/chainreactors/aiscan/pkg/commands" + "github.com/chainreactors/aiscan/pkg/runner" "github.com/chainreactors/aiscan/pkg/tui" "github.com/chainreactors/aiscan/pkg/webproto" scantool "github.com/chainreactors/aiscan/tools/scan" @@ -464,7 +464,8 @@ func (s *Service) runScanViaAgent(ctx context.Context, job *ScanJob) { s.mu.Lock() s.scanAgents[job.ID] = agent.id s.mu.Unlock() - if ctx.Err() != nil { + if err := ctx.Err(); err != nil { + s.finishScanContext(job, err) return } @@ -808,376 +809,6 @@ func (w *sseStreamWriter) Write(p []byte) (int, error) { return len(p), nil } -// defaultReportLang is the language the report is frozen in at scan time; the -// stored copy is only a fallback — GetReport re-renders per request language. -const defaultReportLang = "zh" - -// reportWriter holds the target language so every helper can call w.tr() -// without threading `lang` through every signature. -type reportWriter struct { - strings.Builder - lang string -} - -func newReportWriter(lang string) *reportWriter { - if strings.HasPrefix(strings.ToLower(lang), "zh") { - return &reportWriter{lang: "zh"} - } - return &reportWriter{lang: "en"} -} - -func (w *reportWriter) tr(zh, en string) string { - if w.lang == "zh" { - return zh - } - return en -} - -func (w *reportWriter) modeName(mode string) string { - if strings.EqualFold(mode, "full") { - return w.tr("全面侦察", "Full recon") - } - return w.tr("快速侦察", "Quick recon") -} - -func (w *reportWriter) sep() string { return w.tr("：", ": ") } - -// buildMarkdownReport renders a scan result as an operator-facing recon report. -// It reads like something a human wrote — a prose overview instead of a raw -// metric dump, no internal scanner names (gogo_portscan / check) leaking into -// the prose, and bare live hosts (an icmp echo, say) folded into a trailing -// list rather than each claiming a full section. -func buildMarkdownReport(target, mode string, result *output.Result, lang string) string { - w := newReportWriter(lang) - - heading := output.FirstNonEmpty(target, w.tr("目标", "target")) - fmt.Fprintf(w, "# %s%s\n\n", w.tr("侦察报告 · ", "Recon report · "), heading) - fmt.Fprintf(w, "%s `%s` · %s · %s\n\n", - w.tr("目标", "Target"), target, - w.modeName(mode), - time.Now().Format("2006-01-02 15:04:05")) - w.WriteString("---\n\n") - - if result == nil { - w.WriteString(w.tr("本次扫描未返回结构化结果。\n", "No structured result was returned.\n")) - return w.String() - } - - w.WriteString("## " + w.tr("概述", "Overview") + "\n\n") - w.writeOverview(result) - w.WriteString("\n\n") - - rich, bare := splitReportAssets(result.Assets) - if len(rich) > 0 { - w.WriteString("## " + w.tr("资产明细", "Assets") + "\n\n") - for _, asset := range rich { - w.writeAsset(asset) - } - } - if len(bare) > 0 { - w.WriteString("## " + w.tr("其他存活主机", "Other live hosts") + "\n\n") - for _, asset := range bare { - w.writeBareAsset(asset) - } - w.WriteString("\n") - } - - return w.String() -} - -// writeOverview appends the executive summary — one flowing paragraph that -// names only the numbers that are actually present, so a clean scan reads like -// a sentence rather than a table full of zeros. -func (w *reportWriter) writeOverview(result *output.Result) { - s := result.Summary - hosts := reportHostCount(result.Assets) - fingers := resultFingerprintCount(result) - - if w.lang == "zh" { - fmt.Fprintf(w, "本次侦察共识别 %d 台主机、%d 个开放服务", hosts, s.Services) - if s.Webs > 0 { - fmt.Fprintf(w, "（含 %d 个 Web 站点）", s.Webs) - } - w.WriteString("。") - if s.Probes > 0 { - fmt.Fprintf(w, "累计探测 %d 条路径", s.Probes) - if fingers > 0 { - fmt.Fprintf(w, "、命中 %d 项 Web 指纹", fingers) - } - w.WriteString("。") - } else if fingers > 0 { - fmt.Fprintf(w, "命中 %d 项 Web 指纹。", fingers) - } - if s.Loots > 0 { - fmt.Fprintf(w, "**发现 %d 项需优先复核的安全发现（凭证 / 弱口令 / 漏洞）。**", s.Loots) - } - if s.Errors > 0 { - fmt.Fprintf(w, "另有 %d 处探测报错。", s.Errors) - } - if s.Duration != "" { - fmt.Fprintf(w, "全程耗时 %s。", s.Duration) - } - return - } - - fmt.Fprintf(w, "The scan identified %s across %s", plural(hosts, "host", "hosts"), plural(s.Services, "open service", "open services")) - if s.Webs > 0 { - fmt.Fprintf(w, " (%s)", plural(s.Webs, "web site", "web sites")) - } - w.WriteString(". ") - if s.Probes > 0 { - fmt.Fprintf(w, "It probed %s", plural(s.Probes, "path", "paths")) - if fingers > 0 { - fmt.Fprintf(w, " and matched %s", plural(fingers, "fingerprint", "fingerprints")) - } - w.WriteString(". ") - } else if fingers > 0 { - fmt.Fprintf(w, "It matched %s. ", plural(fingers, "fingerprint", "fingerprints")) - } - if s.Loots > 0 { - fmt.Fprintf(w, "**%s surfaced (credentials / weak passwords / vulnerabilities) — review these first.** ", plural(s.Loots, "security finding", "security findings")) - } - if s.Errors > 0 { - fmt.Fprintf(w, "%s occurred during probing. ", plural(s.Errors, "error", "errors")) - } - if s.Duration != "" { - fmt.Fprintf(w, "The scan took %s.", s.Duration) - } -} - -func plural(n int, one, many string) string { - if n == 1 { - return fmt.Sprintf("%d %s", n, one) - } - return fmt.Sprintf("%d %s", n, many) -} - -// reportHostCount collapses assets down to distinct hosts, so an IP that has -// both an icmp echo and a web service counts once, not twice. -func reportHostCount(assets []output.Asset) int { - seen := make(map[string]struct{}) - for _, a := range assets { - if h := assetHost(a); h != "" { - seen[h] = struct{}{} - } - } - if len(seen) == 0 { - return len(assets) - } - return len(seen) -} - -func assetHost(a output.Asset) string { - v := output.FirstNonEmpty(a.Target, a.Key, a.Title) - if i := strings.Index(v, "://"); i >= 0 { - v = v[i+3:] - } - if i := strings.IndexAny(v, "/?#"); i >= 0 { - v = v[:i] - } - if strings.Count(v, ":") == 1 { // host:port — drop the port, leave IPv6 alone - v = v[:strings.LastIndex(v, ":")] - } - return v -} - -func splitReportAssets(assets []output.Asset) (rich, bare []output.Asset) { - for _, a := range assets { - if assetIsBare(a) { - bare = append(bare, a) - } else { - rich = append(rich, a) - } - } - return rich, bare -} - -// assetIsBare is true for a live host that only answered with non-web services -// (an icmp echo, a bare tcp port) — nothing worth its own section. -func assetIsBare(a output.Asset) bool { - hasService := false - for _, item := range a.Items { - if item.Kind != output.AssetItemService { - return false - } - hasService = true - svc := strings.ToLower(output.AssetDataString(item.Data, "service") + " " + output.AssetDataString(item.Data, "protocol")) - if strings.Contains(svc, "http") { - return false - } - } - return hasService -} - -func (w *reportWriter) writeAsset(asset output.Asset) { - title := output.FirstNonEmpty(asset.Title, asset.Target, asset.Key, w.tr("资产", "Asset")) - if asset.Target != "" && asset.Target != title { - fmt.Fprintf(w, "### %s — `%s`\n\n", title, asset.Target) - } else { - fmt.Fprintf(w, "### %s\n\n", title) - } - - w.writeFact(w.tr("开放服务", "Services"), assetServiceFacts(asset.Items)) - w.writeFact(w.tr("HTTP 响应", "HTTP"), assetHTTPStatuses(asset.Items)) - w.writeFact(w.tr("Web 指纹", "Fingerprints"), assetFingers(asset.Items)) - if paths := assetPathCount(asset.Items); paths > 0 { - fmt.Fprintf(w, "- %s%s%s\n", w.tr("已探测路径", "Paths"), w.sep(), w.tr(fmt.Sprintf("%d 条", paths), fmt.Sprintf("%d", paths))) - } - if asset.Status != "" { - fmt.Fprintf(w, "- %s%s%s\n", w.tr("状态", "State"), w.sep(), markdownCode(asset.Status)) - } - w.WriteString("\n") - - w.writeLootMarkdown(asset.Items) -} - -func (w *reportWriter) writeBareAsset(asset output.Asset) { - host := output.FirstNonEmpty(asset.Target, asset.Title, asset.Key) - if services := assetServiceFacts(asset.Items); len(services) > 0 { - fmt.Fprintf(w, "- `%s` · %s\n", host, strings.Join(services, ", ")) - } else { - fmt.Fprintf(w, "- `%s`\n", host) - } -} - -func (w *reportWriter) writeFact(label string, values []string) { - if len(values) == 0 { - return - } - coded := make([]string, 0, len(values)) - for _, value := range values { - coded = append(coded, markdownCode(value)) - } - fmt.Fprintf(w, "- %s%s%s\n", label, w.sep(), strings.Join(coded, w.tr("、", ", "))) -} - -func (w *reportWriter) writeLootMarkdown(items []output.AssetItem) { - wrote := false - for _, item := range items { - switch item.Kind { - case output.AssetItemLoot, output.AssetItemNote, output.AssetItemResponse, output.AssetItemError: - summary := output.FirstNonEmpty(item.Summary, item.Title) - detail := output.AssetItemDetail(item) - if summary == "" && detail == "" { - continue - } - if !wrote { - w.WriteString("#### " + w.tr("分析研判", "Analysis") + "\n\n") - wrote = true - } - if summary == "" { - summary = firstMarkdownLine(detail) - } - fmt.Fprintf(w, "##### %s\n\n", markdownHeading(summary)) - if detail != "" && !sameMarkdownText(summary, detail) { - writeMarkdownBlock(&w.Builder, detail) - } else if detail == "" && summary != "" { - w.WriteString(summary) - w.WriteString("\n\n") - } - } - } -} - -func firstMarkdownLine(value string) string { - value = strings.TrimSpace(value) - if value == "" { - return "" - } - if idx := strings.IndexByte(value, '\n'); idx >= 0 { - return strings.TrimSpace(value[:idx]) - } - return value -} - -func sameMarkdownText(left, right string) bool { - return strings.TrimSpace(left) == strings.TrimSpace(right) -} - -func writeMarkdownBlock(sb *strings.Builder, value string) { - value = strings.TrimSpace(value) - if value == "" { - return - } - sb.WriteString(value) - sb.WriteString("\n\n") -} - -func assetServiceFacts(items []output.AssetItem) []string { - var values []string - for _, item := range items { - if item.Kind != output.AssetItemService { - continue - } - values = append(values, strings.Join(output.CompactStrings( - output.AssetDataString(item.Data, "protocol"), - output.AssetDataString(item.Data, "service"), - output.AssetDataString(item.Data, "port"), - ), " ")) - } - return output.CompactStrings(values...) -} - -func assetHTTPStatuses(items []output.AssetItem) []string { - var values []string - for _, item := range items { - if item.Kind == output.AssetItemPath && item.Status != "" { - values = append(values, item.Status) - } - } - return output.CompactStrings(values...) -} - -func assetFingers(items []output.AssetItem) []string { - var values []string - for _, item := range items { - switch item.Kind { - case output.AssetItemFingerprint: - values = append(values, output.FirstNonEmpty(item.Title, output.AssetDataString(item.Data, "name"))) - case output.AssetItemPath: - values = append(values, output.AssetDataStrings(item.Data, "fingers")...) - } - } - return output.CompactStrings(values...) -} - -func assetPathCount(items []output.AssetItem) int { - count := 0 - for _, item := range items { - if item.Kind == output.AssetItemPath { - count++ - } - } - return count -} - -func resultFingerprintCount(result *output.Result) int { - if result == nil { - return 0 - } - seen := make(map[string]struct{}) - for _, asset := range result.Assets { - for _, finger := range assetFingers(asset.Items) { - seen[strings.ToLower(finger)] = struct{}{} - } - } - return len(seen) -} - -func markdownCode(value string) string { - value = strings.ReplaceAll(value, "`", "'") - return "`" + value + "`" -} - -func markdownHeading(value string) string { - value = strings.TrimSpace(value) - value = strings.ReplaceAll(value, "\n", " ") - if value == "" { - return "Analysis" - } - return strings.TrimLeft(value, "# ") -} - func generateID() string { b := make([]byte, 16) _, _ = rand.Read(b) diff --git a/pkg/web/sse_test.go b/pkg/web/sse_test.go index 8ff158e8..4c8724b1 100644 --- a/pkg/web/sse_test.go +++ b/pkg/web/sse_test.go @@ -9,9 +9,9 @@ import ( "testing" "time" + "github.com/chainreactors/aiscan/core/aop" + xeval "github.com/chainreactors/aiscan/core/aop/x/eval" "github.com/chainreactors/aiscan/core/output" - "github.com/chainreactors/aiscan/pkg/aop" - xeval "github.com/chainreactors/aiscan/pkg/aop/x/eval" ) // A saturated subscriber buffer must never swallow a reliable terminal event. diff --git a/pkg/web/store_sqlite.go b/pkg/web/store_sqlite.go index f9e55441..bdb429f9 100644 --- a/pkg/web/store_sqlite.go +++ b/pkg/web/store_sqlite.go @@ -8,8 +8,8 @@ import ( "strings" "time" + "github.com/chainreactors/aiscan/core/aop" "github.com/chainreactors/aiscan/core/output" - "github.com/chainreactors/aiscan/pkg/aop" "github.com/chainreactors/aiscan/pkg/webproto" _ "modernc.org/sqlite" ) diff --git a/pkg/web/store_sqlite_test.go b/pkg/web/store_sqlite_test.go index 714d6427..fae1760a 100644 --- a/pkg/web/store_sqlite_test.go +++ b/pkg/web/store_sqlite_test.go @@ -8,8 +8,8 @@ import ( "testing" "time" + "github.com/chainreactors/aiscan/core/aop" "github.com/chainreactors/aiscan/core/output" - "github.com/chainreactors/aiscan/pkg/aop" ) func createStoredSession(t *testing.T, store *SQLiteStore, id string) { diff --git a/pkg/web/types.go b/pkg/web/types.go index b6206301..58618590 100644 --- a/pkg/web/types.go +++ b/pkg/web/types.go @@ -5,8 +5,8 @@ import ( "errors" "time" + "github.com/chainreactors/aiscan/core/aop" "github.com/chainreactors/aiscan/core/output" - "github.com/chainreactors/aiscan/pkg/aop" "github.com/chainreactors/aiscan/pkg/webproto" ) diff --git a/pkg/webagent/agent.go b/pkg/webagent/agent.go index 142c920d..9244c832 100644 --- a/pkg/webagent/agent.go +++ b/pkg/webagent/agent.go @@ -9,11 +9,11 @@ import ( "path/filepath" "strings" + "github.com/chainreactors/aiscan/agent" + "github.com/chainreactors/aiscan/core/aop" cfg "github.com/chainreactors/aiscan/core/config" - "github.com/chainreactors/aiscan/core/runner" - "github.com/chainreactors/aiscan/pkg/agent" - "github.com/chainreactors/aiscan/pkg/aop" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/core/telemetry" + "github.com/chainreactors/aiscan/pkg/runner" "github.com/chainreactors/aiscan/pkg/webproto" "github.com/chainreactors/ioa/protocols" "github.com/chainreactors/utils/pty" @@ -37,7 +37,7 @@ func RunWebSocket(ctx context.Context, option *cfg.Option, logger telemetry.Logg return err } - appConfig := cfg.AppConfig(option, cfg.RuntimeFeatures{ + appConfig := runner.AppConfig(option, runner.RuntimeFeatures{ ProviderEnabled: true, ProviderOptional: true, ToolsEnabled: true, AIEnabled: true, }, logger) appConfig.IOA = remoteIOAConfig(option, identityRef) @@ -46,7 +46,7 @@ func RunWebSocket(ctx context.Context, option *cfg.Option, logger telemetry.Logg return err } defer application.Close() - cfg.ApplyResolvedProviderOptions(option, application.ProviderConfig) + runner.ApplyResolvedProviderOptions(option, application.ProviderConfig) rt, err := runner.NewAgentRuntime(ctx, option, logger, &runner.RuntimeConfig{ ExistingApp: application, NoOutput: true, REPLMode: runner.REPLPersistent, }) @@ -172,7 +172,7 @@ func reloadAgentConfig(serverURL string, rt *runner.AgentRuntime, app *runner.Ap logger.Warnf("config reload: fetch remote config: %s", err) return nil, "", err } - providerConfig := cfg.ProviderConfig(remoteOpt) + providerConfig := runner.ProviderConfig(remoteOpt) resolved, err := agent.ResolveProvider(&providerConfig) if err != nil { logger.Warnf("config reload: resolve provider: %s", err) @@ -321,11 +321,11 @@ func webNodeRef(option *cfg.Option) (protocols.NodeRef, error) { return protocols.NodeRef{ID: name, Authority: authority}, nil } -func remoteIOAConfig(option *cfg.Option, ref protocols.NodeRef) *cfg.IOAConfig { +func remoteIOAConfig(option *cfg.Option, ref protocols.NodeRef) *runner.IOAConfig { if option == nil || option.IOAURL == "" { return nil } - return &cfg.IOAConfig{ + return &runner.IOAConfig{ URL: option.IOAURL, NodeID: option.IOANodeID, NodeName: option.IOANodeName, diff --git a/pkg/webagent/agent_test.go b/pkg/webagent/agent_test.go index 630acc9b..1c550f6c 100644 --- a/pkg/webagent/agent_test.go +++ b/pkg/webagent/agent_test.go @@ -13,10 +13,11 @@ import ( "testing" "time" + "github.com/chainreactors/aiscan/core/aop" + "github.com/chainreactors/aiscan/core/capability" cfg "github.com/chainreactors/aiscan/core/config" "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" - "github.com/chainreactors/aiscan/pkg/aop" "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/aiscan/pkg/webproto" "github.com/chainreactors/ioa/protocols" @@ -346,7 +347,7 @@ func TestRunConnectionPTYRoundTrip(t *testing.T) { defer cancel() reg := commands.NewRegistry() - commands.BuildGroup("core", &commands.Deps{WorkDir: t.TempDir(), BashTimeout: 5}, reg) + commands.BuildPlan(capability.Select(capability.Options{Groups: []string{"core"}}), &commands.Deps{WorkDir: t.TempDir(), BashTimeout: 5}, reg) done := make(chan error, 1) go func() { @@ -427,7 +428,7 @@ func TestRunConnectionPushesPTYSessionsOnManagerEvents(t *testing.T) { defer cancel() reg := commands.NewRegistry() - commands.BuildGroup("core", &commands.Deps{WorkDir: t.TempDir(), BashTimeout: 5}, reg) + commands.BuildPlan(capability.Select(capability.Options{Groups: []string{"core"}}), &commands.Deps{WorkDir: t.TempDir(), BashTimeout: 5}, reg) mgr := RegistryPTYManager(reg) if mgr == nil { t.Fatal("bash command did not expose tmux manager") diff --git a/pkg/webagent/aop_tool.go b/pkg/webagent/aop_tool.go index 75511a33..9a6ff8ad 100644 --- a/pkg/webagent/aop_tool.go +++ b/pkg/webagent/aop_tool.go @@ -8,10 +8,10 @@ import ( "strings" "time" + "github.com/chainreactors/aiscan/core/aop" "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" "github.com/chainreactors/aiscan/core/tool" - "github.com/chainreactors/aiscan/pkg/aop" "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/aiscan/pkg/webproto" ) diff --git a/pkg/webagent/aop_tool_test.go b/pkg/webagent/aop_tool_test.go index dee00af7..7c79ab25 100644 --- a/pkg/webagent/aop_tool_test.go +++ b/pkg/webagent/aop_tool_test.go @@ -7,10 +7,10 @@ import ( "testing" "time" + "github.com/chainreactors/aiscan/core/aop" "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" "github.com/chainreactors/aiscan/core/tool" - "github.com/chainreactors/aiscan/pkg/aop" "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/aiscan/pkg/webproto" ) diff --git a/pkg/webagent/connection.go b/pkg/webagent/connection.go index 071d96e1..88888bca 100644 --- a/pkg/webagent/connection.go +++ b/pkg/webagent/connection.go @@ -8,12 +8,12 @@ import ( "sync" "time" + "github.com/chainreactors/aiscan/agent" + "github.com/chainreactors/aiscan/core/aop" "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" - "github.com/chainreactors/aiscan/pkg/agent" - "github.com/chainreactors/aiscan/pkg/aop" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" "github.com/chainreactors/aiscan/pkg/webproto" "github.com/chainreactors/ioa/protocols" "github.com/chainreactors/utils/pty" diff --git a/pkg/webagent/connection_lifecycle_test.go b/pkg/webagent/connection_lifecycle_test.go index 2d908215..a5233f4c 100644 --- a/pkg/webagent/connection_lifecycle_test.go +++ b/pkg/webagent/connection_lifecycle_test.go @@ -9,8 +9,8 @@ import ( "testing" "time" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" "github.com/chainreactors/aiscan/pkg/webproto" "github.com/chainreactors/ioa/protocols" "github.com/gorilla/websocket" diff --git a/pkg/webagent/pty.go b/pkg/webagent/pty.go index b13c678b..a7b557e2 100644 --- a/pkg/webagent/pty.go +++ b/pkg/webagent/pty.go @@ -5,7 +5,7 @@ import ( "sync" "time" - "github.com/chainreactors/aiscan/pkg/agent/tmux" + "github.com/chainreactors/aiscan/agent/tmux" "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/aiscan/pkg/webproto" "github.com/chainreactors/utils/pty" diff --git a/pkg/webagent/stream.go b/pkg/webagent/stream.go index 5d553f3f..1e091fc1 100644 --- a/pkg/webagent/stream.go +++ b/pkg/webagent/stream.go @@ -3,7 +3,7 @@ package webagent import ( "sync" - "github.com/chainreactors/aiscan/pkg/aop" + "github.com/chainreactors/aiscan/core/aop" "github.com/chainreactors/aiscan/pkg/webproto" ) diff --git a/pkg/webagent/toolnode.go b/pkg/webagent/toolnode.go index a2fc7e18..23a6a7de 100644 --- a/pkg/webagent/toolnode.go +++ b/pkg/webagent/toolnode.go @@ -9,8 +9,8 @@ import ( "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" "github.com/chainreactors/aiscan/pkg/webproto" "github.com/chainreactors/ioa/protocols" ) diff --git a/pkg/webagent/toolnode_test.go b/pkg/webagent/toolnode_test.go index 7a7cb120..9fbca5ca 100644 --- a/pkg/webagent/toolnode_test.go +++ b/pkg/webagent/toolnode_test.go @@ -14,9 +14,9 @@ import ( "github.com/gorilla/websocket" + "github.com/chainreactors/aiscan/core/aop" "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" - "github.com/chainreactors/aiscan/pkg/aop" "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/aiscan/pkg/webproto" ) diff --git a/pkg/webproto/message.go b/pkg/webproto/message.go index e1225fb2..71a60261 100644 --- a/pkg/webproto/message.go +++ b/pkg/webproto/message.go @@ -5,7 +5,7 @@ import ( "fmt" "github.com/chainreactors/aiscan/core/tool" - "github.com/chainreactors/aiscan/pkg/aop" + "github.com/chainreactors/aiscan/core/aop" "github.com/chainreactors/ioa/protocols" "github.com/chainreactors/utils/pty" ) diff --git a/pkg/webproto/message_test.go b/pkg/webproto/message_test.go index 799759e5..320fa7f8 100644 --- a/pkg/webproto/message_test.go +++ b/pkg/webproto/message_test.go @@ -6,7 +6,7 @@ import ( "strings" "testing" - "github.com/chainreactors/aiscan/pkg/aop" + "github.com/chainreactors/aiscan/core/aop" "github.com/chainreactors/utils/pty" ) diff --git a/skills/availability.go b/skills/availability.go index 68aaa763..e9711891 100644 --- a/skills/availability.go +++ b/skills/availability.go @@ -1,15 +1,7 @@ package skills -var blocked = map[string]bool{ - "katana": true, - "passive": true, -} - -//nolint:unused // called from build-tagged files -func enableSkill(name string) { - delete(blocked, name) -} +import "github.com/chainreactors/aiscan/core/capability" func skillAvailable(name string) bool { - return !blocked[name] + return capability.SkillEnabled(name) } diff --git a/skills/availability_full.go b/skills/availability_full.go deleted file mode 100644 index d8096464..00000000 --- a/skills/availability_full.go +++ /dev/null @@ -1,8 +0,0 @@ -//go:build full - -package skills - -func init() { - enableSkill("katana") - enableSkill("passive") -} diff --git a/test-skips.json b/test-skips.json new file mode 100644 index 00000000..aaa353fd --- /dev/null +++ b/test-skips.json @@ -0,0 +1,198 @@ +[ + { + "path": "agent/agent_test.go", + "format": "unix-only test", + "count": 5, + "category": "platform", + "reason": "Exercises POSIX process, signal, or shell semantics and is covered by Linux CI." + }, + { + "path": "agent/agent_test.go", + "format": "python3 not found", + "count": 1, + "category": "external_runtime", + "reason": "The subagent process fixture requires a locally installed Python 3 interpreter." + }, + { + "path": "agent/agent_test.go", + "format": "no LIVE_TEST_API_KEY set; skipping live LLM test", + "count": 1, + "category": "live_llm", + "reason": "Calls a live model and therefore requires an explicitly supplied credential." + }, + { + "path": "agent/helpers_test.go", + "format": "set TEST_API_KEY, TEST_BASE_URL, TEST_MODEL to run live tests", + "count": 1, + "category": "live_llm", + "reason": "Provider conformance against a live endpoint requires explicit credentials and model selection." + }, + { + "path": "agent/provider/cache_test.go", + "format": "set TEST_API_KEY, TEST_BASE_URL, TEST_MODEL to run live cache test", + "count": 1, + "category": "live_llm", + "reason": "Prompt-cache behavior can only be verified against a credentialed live provider." + }, + { + "path": "agent/provider/cache_test.go", + "format": "%s: provider does not support streaming", + "count": 1, + "category": "capability", + "reason": "The shared provider matrix includes implementations that intentionally do not expose streaming." + }, + { + "path": "agent/provider/cache_test.go", + "format": "set TEST_API_KEY, TEST_BASE_URL, TEST_MODEL", + "count": 1, + "category": "live_llm", + "reason": "Live multi-provider cache coverage requires explicit endpoint credentials." + }, + { + "path": "agent/tmux/manager_test.go", + "format": "unix-only test", + "count": 19, + "category": "platform", + "reason": "Validates PTY, process-group, and signal behavior provided only by the Unix tmux implementation." + }, + { + "path": "agent/tmux/manager_test.go", + "format": "unix-only", + "count": 2, + "category": "platform", + "reason": "Validates Unix-only PTY and process lifecycle behavior." + }, + { + "path": "pkg/commands/bash_test.go", + "format": "unix-only test", + "count": 2, + "category": "platform", + "reason": "Exercises POSIX shell execution details covered by Linux CI." + }, + { + "path": "pkg/commands/bash_test.go", + "format": "unix-only", + "count": 8, + "category": "platform", + "reason": "Exercises POSIX shell, process, or permission semantics." + }, + { + "path": "pkg/commands/bash_test.go", + "format": "shell assertions are unix-only", + "count": 5, + "category": "platform", + "reason": "Assertions depend on POSIX shell syntax and command behavior." + }, + { + "path": "pkg/commands/tmux_test.go", + "format": "unix-only", + "count": 14, + "category": "platform", + "reason": "The tmux command integration requires Unix PTY and signal semantics." + }, + { + "path": "pkg/headless/engine_test.go", + "format": "skip: requires runtime DSL variables", + "count": 1, + "category": "capability", + "reason": "Two mixed HTTP/headless templates reference variables unavailable to the headless-only compile fixture." + }, + { + "path": "pkg/tui/console_test.go", + "format": "shell assertion is unix-only", + "count": 1, + "category": "platform", + "reason": "The assertion targets POSIX shell rendering." + }, + { + "path": "pkg/web/agents_test.go", + "format": "chromium not found, skipping browser e2e test", + "count": 1, + "category": "external_runtime", + "reason": "Browser E2E requires a Chromium-family executable; CI installs and exercises it." + }, + { + "path": "tools/arsenal/arsenal_tool_test.go", + "format": "skip network test in short mode", + "count": 1, + "category": "external_api", + "reason": "The case performs a real network fetch and is intentionally excluded by go test -short." + }, + { + "path": "tools/arsenal/arsenal_tool_test.go", + "format": "e2e only on linux/amd64", + "count": 1, + "category": "platform", + "reason": "The external arsenal binary fixture is published only for linux/amd64." + }, + { + "path": "tools/functional_integration_full_test.go", + "format": "set AISCAN_INTEGRATION=1 to run public network regression tests", + "count": 1, + "category": "external_api", + "reason": "The scheduled full scanner regression contacts public network targets." + }, + { + "path": "tools/functional_integration_test.go", + "format": "set AISCAN_INTEGRATION=1 to run public network regression tests", + "count": 1, + "category": "external_api", + "reason": "The scheduled scanner regression contacts public network targets." + }, + { + "path": "tools/ioa/commands_test.go", + "format": "set LIVE_TEST_API_KEY or DEEPSEEK_API_KEY to run live LLM IOA test", + "count": 1, + "category": "live_llm", + "reason": "The IOA integration calls a live model endpoint and requires an explicit credential." + }, + { + "path": "tools/playwright/browser_test.go", + "format": "no Chromium/Chrome found, skipping browser integration test", + "count": 1, + "category": "external_runtime", + "reason": "The browser integration requires a locally installed Chromium-family executable." + }, + { + "path": "tools/playwright/recorder_test.go", + "format": "no Chromium/Chrome found, skipping browser integration test", + "count": 1, + "category": "external_runtime", + "reason": "The recorder integration requires a locally installed Chromium-family executable." + }, + { + "path": "tools/proton/command_test.go", + "format": "unix-only", + "count": 1, + "category": "platform", + "reason": "The command fixture asserts POSIX file permission semantics." + }, + { + "path": "tools/register_command_integration_test.go", + "format": "set AISCAN_INTEGRATION=1 to run", + "count": 2, + "category": "external_api", + "reason": "The passive search integrations are opt-in because they call external services." + }, + { + "path": "tools/register_command_integration_test.go", + "format": "FOFA_EMAIL / FOFA_KEY required", + "count": 1, + "category": "external_api", + "reason": "The FOFA integration requires user-supplied service credentials." + }, + { + "path": "tools/register_command_integration_test.go", + "format": "HUNTER_TOKEN or HUNTER_API_KEY required", + "count": 1, + "category": "external_api", + "reason": "The Hunter integration requires a user-supplied service credential." + }, + { + "path": "web/frontend/e2e/aiscan-web.spec.ts", + "format": "LLM_API_KEY env var required", + "count": 2, + "category": "live_llm", + "reason": "Provider connectivity and model discovery checks are retained as opt-in live endpoint coverage." + } +] diff --git a/tools/arsenal/register.go b/tools/arsenal/register.go index e0b8a550..59c0613e 100644 --- a/tools/arsenal/register.go +++ b/tools/arsenal/register.go @@ -1,10 +1,14 @@ package arsenal -import "github.com/chainreactors/aiscan/pkg/commands" +import ( + "github.com/chainreactors/aiscan/core/capability" + "github.com/chainreactors/aiscan/pkg/commands" +) func init() { + capability.Register(capability.Descriptor{ID: "arsenal", Kind: capability.KindTool, Group: "arsenal"}) commands.RegisterFactory(commands.Factory{ - Group: "arsenal", + Capability: "arsenal", Build: func(deps *commands.Deps, reg *commands.CommandRegistry) { logger := deps.GetLogger() diff --git a/tools/capability_test.go b/tools/capability_test.go new file mode 100644 index 00000000..4db178aa --- /dev/null +++ b/tools/capability_test.go @@ -0,0 +1,10 @@ +package tools + +import ( + "github.com/chainreactors/aiscan/core/capability" + "github.com/chainreactors/aiscan/pkg/commands" +) + +func buildTestGroups(groups []string, deps *commands.Deps, reg *commands.CommandRegistry) { + commands.BuildPlan(capability.Select(capability.Options{Groups: groups}), deps, reg) +} diff --git a/tools/functional_integration_full_test.go b/tools/functional_integration_full_test.go index e84ab212..9dfa19d9 100644 --- a/tools/functional_integration_full_test.go +++ b/tools/functional_integration_full_test.go @@ -9,10 +9,11 @@ import ( "testing" "time" + "github.com/chainreactors/aiscan/core/capability" "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" _ "github.com/chainreactors/aiscan/tools/katana" "github.com/chainreactors/aiscan/tools/scan/engine" ) @@ -25,9 +26,9 @@ func TestFullScannerPublicIntegration(t *testing.T) { bus := eventbus.New[output.ToolDataEvent]() recorder := newFunctionalRecorder(bus) registry := commands.NewRegistry() - commands.BuildGroup("scanner", &commands.Deps{ - WorkDir: t.TempDir(), EngineSet: &engine.Set{}, DataBus: bus, Logger: telemetry.NopLogger(), - }, registry) + deps := &commands.Deps{WorkDir: t.TempDir(), DataBus: bus, Logger: telemetry.NopLogger()} + commands.Provide(deps, engine.SetKey, &engine.Set{}) + commands.BuildPlan(capability.Select(capability.Options{Groups: []string{"scanner"}}), deps, registry) runFunctionalCases(t, registry, recorder, []functionalCase{{ Name: "katana/redhaze-depth-one", Tool: "katana", diff --git a/tools/functional_integration_test.go b/tools/functional_integration_test.go index 3a415a16..3dbe9689 100644 --- a/tools/functional_integration_test.go +++ b/tools/functional_integration_test.go @@ -10,11 +10,12 @@ import ( "testing" "time" + "github.com/chainreactors/aiscan/core/capability" "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" "github.com/chainreactors/aiscan/core/resources" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" _ "github.com/chainreactors/aiscan/tools/gogo" _ "github.com/chainreactors/aiscan/tools/neutron" "github.com/chainreactors/aiscan/tools/scan/engine" @@ -38,10 +39,13 @@ func TestScannerPublicIntegration(t *testing.T) { bus := eventbus.New[output.ToolDataEvent]() recorder := newFunctionalRecorder(bus) registry := commands.NewRegistry() - commands.BuildGroup("scanner", &commands.Deps{ - WorkDir: t.TempDir(), EngineSet: engineSet, Resources: engineSet.Resources, + deps := &commands.Deps{ + WorkDir: t.TempDir(), DataBus: bus, Logger: telemetry.NopLogger(), - }, registry) + } + commands.Provide(deps, engine.SetKey, engineSet) + commands.Provide(deps, resources.SetKey, engineSet.Resources) + commands.BuildPlan(capability.Select(capability.Options{Groups: []string{"scanner"}}), deps, registry) templateFile := filepath.Join(t.TempDir(), "redhaze-marker.yaml") writeTestFile(t, templateFile, `id: redhaze-public-marker info: diff --git a/tools/functional_norace_test.go b/tools/functional_norace_test.go deleted file mode 100644 index 7deb8dd4..00000000 --- a/tools/functional_norace_test.go +++ /dev/null @@ -1,5 +0,0 @@ -//go:build !race - -package tools - -const functionalRaceEnabled = false diff --git a/tools/functional_race_test.go b/tools/functional_race_test.go deleted file mode 100644 index 424133b0..00000000 --- a/tools/functional_race_test.go +++ /dev/null @@ -1,5 +0,0 @@ -//go:build race - -package tools - -const functionalRaceEnabled = true diff --git a/tools/functional_regression_full_test.go b/tools/functional_regression_full_test.go index d1457a96..fa05cc1b 100644 --- a/tools/functional_regression_full_test.go +++ b/tools/functional_regression_full_test.go @@ -9,10 +9,11 @@ import ( "testing" "time" + "github.com/chainreactors/aiscan/core/capability" "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" _ "github.com/chainreactors/aiscan/tools/katana" passivecmd "github.com/chainreactors/aiscan/tools/passive" "github.com/chainreactors/aiscan/tools/scan/engine" @@ -25,12 +26,13 @@ func TestFullScannerFunctionalRegression(t *testing.T) { recorder := newFunctionalRecorder(bus) registry := commands.NewRegistry() engineSet := &engine.Set{} - commands.BuildGroup("scanner", &commands.Deps{ - WorkDir: t.TempDir(), - EngineSet: engineSet, - DataBus: bus, - Logger: telemetry.NopLogger(), - }, registry) + deps := &commands.Deps{ + WorkDir: t.TempDir(), + DataBus: bus, + Logger: telemetry.NopLogger(), + } + commands.Provide(deps, engine.SetKey, engineSet) + commands.BuildPlan(capability.Select(capability.Options{Groups: []string{"scanner"}}), deps, registry) passiveEngine := &functionalPassiveEngine{} passive := passivecmd.New(passiveEngine).WithLogger(telemetry.NopLogger()) diff --git a/tools/functional_regression_test.go b/tools/functional_regression_test.go index 3547584c..5d62fe1d 100644 --- a/tools/functional_regression_test.go +++ b/tools/functional_regression_test.go @@ -16,11 +16,12 @@ import ( "testing" "time" + "github.com/chainreactors/aiscan/core/capability" "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" "github.com/chainreactors/aiscan/core/resources" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" _ "github.com/chainreactors/aiscan/tools/gogo" _ "github.com/chainreactors/aiscan/tools/neutron" _ "github.com/chainreactors/aiscan/tools/proton" @@ -59,13 +60,14 @@ func TestScannerFunctionalRegression(t *testing.T) { bus := eventbus.New[output.ToolDataEvent]() recorder := newFunctionalRecorder(bus) registry := commands.NewRegistry() - commands.BuildGroup("scanner", &commands.Deps{ - WorkDir: workDir, - EngineSet: engineSet, - Resources: engineSet.Resources, - DataBus: bus, - Logger: telemetry.NopLogger(), - }, registry) + deps := &commands.Deps{ + WorkDir: workDir, + DataBus: bus, + Logger: telemetry.NopLogger(), + } + commands.Provide(deps, engine.SetKey, engineSet) + commands.Provide(deps, resources.SetKey, engineSet.Resources) + commands.BuildPlan(capability.Select(capability.Options{Groups: []string{"scanner"}}), deps, registry) required := []string{"scan", "gogo", "spray", "zombie", "neutron", "proton"} for _, name := range required { @@ -97,8 +99,7 @@ http: cases := []functionalCase{ { Name: "gogo/http-fingerprint-jsonl", Tool: "gogo", - Args: []string{"-i", host, "-p", port, "-v", "-o", "jl", "-t", "20"}, - SkipUnderRace: true, + Args: []string{"-i", host, "-p", port, "-v", "-o", "jl", "-t", "20"}, Check: func(t *testing.T, result functionalResult) { requireOutputContains(t, result, `"port":"`+port+`"`, "nginx") requireEvent(t, result, "gogo", output.ToolDataService, func(data any) bool { @@ -113,8 +114,7 @@ http: }, { Name: "gogo/target-file", Tool: "gogo", - Args: []string{"-l", targetsFile, "-p", port, "-o", "jl", "-t", "20"}, - SkipUnderRace: true, + Args: []string{"-l", targetsFile, "-p", port, "-o", "jl", "-t", "20"}, Check: func(t *testing.T, result functionalResult) { requireOutputContains(t, result, `"port":"`+port+`"`) }, @@ -196,9 +196,8 @@ http: }, { Name: "scan/quick-pipeline", Tool: "scan", - Args: []string{"-i", host, "--ports", port, "--mode", "quick", "--verify=off", "--timeout", "2", "--no-color"}, - Timeout: 30 * time.Second, - SkipUnderRace: true, + Args: []string{"-i", host, "--ports", port, "--mode", "quick", "--verify=off", "--timeout", "2", "--no-color"}, + Timeout: 30 * time.Second, Check: func(t *testing.T, result functionalResult) { requireOutputContains(t, result, "[summary] completed", port) requireEvent(t, result, "gogo", output.ToolDataService, nil) diff --git a/tools/functional_testkit_test.go b/tools/functional_testkit_test.go index d32e0348..e36121a4 100644 --- a/tools/functional_testkit_test.go +++ b/tools/functional_testkit_test.go @@ -22,13 +22,12 @@ type functionalResult struct { } type functionalCase struct { - Name string - Tool string - Args []string - Stdin string - Timeout time.Duration - SkipUnderRace bool - Check func(*testing.T, functionalResult) + Name string + Tool string + Args []string + Stdin string + Timeout time.Duration + Check func(*testing.T, functionalResult) } type functionalRecorder struct { @@ -62,9 +61,6 @@ func runFunctionalCases(t *testing.T, registry *commands.CommandRegistry, record t.Helper() for _, testCase := range cases { t.Run(testCase.Name, func(t *testing.T) { - if functionalRaceEnabled && testCase.SkipUnderRace { - t.Skip("upstream scanner has a known internal race") - } if !registry.Has(testCase.Tool) { t.Fatalf("tool %q is not registered", testCase.Tool) } diff --git a/tools/gogo/gogo.go b/tools/gogo/gogo.go index 71b42381..bfa07f4a 100644 --- a/tools/gogo/gogo.go +++ b/tools/gogo/gogo.go @@ -9,8 +9,8 @@ import ( "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" "github.com/chainreactors/aiscan/tools/toolargs" gogocore "github.com/chainreactors/gogo/v2/core" "github.com/chainreactors/sdk/gogo" diff --git a/tools/gogo/gogo_test.go b/tools/gogo/gogo_test.go index 93dcef73..04e7cef6 100644 --- a/tools/gogo/gogo_test.go +++ b/tools/gogo/gogo_test.go @@ -9,8 +9,8 @@ import ( "sync/atomic" "testing" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" gogopkg "github.com/chainreactors/gogo/v2/pkg" sdkgogo "github.com/chainreactors/sdk/gogo" ) diff --git a/tools/gogo/register.go b/tools/gogo/register.go index 65eb87a7..5869e3a1 100644 --- a/tools/gogo/register.go +++ b/tools/gogo/register.go @@ -1,22 +1,28 @@ package gogo import ( - cfg "github.com/chainreactors/aiscan/core/config" + "github.com/chainreactors/aiscan/core/capability" + "github.com/chainreactors/aiscan/core/deps" "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/aiscan/tools/scan/engine" ) func init() { - cfg.ExtraScannerUsage["gogo"] = func() string { return New(nil).Usage() } + capability.Register(capability.Descriptor{ + ID: "gogo", Kind: capability.KindScanner, Group: "scanner", + CLIName: "gogo", Summary: "gogo", UsageLine: " gogo Run gogo directly", + Usage: func() string { return New(nil).Usage() }, Requires: []string{"scan.engine.Set.Gogo"}, + }) commands.RegisterFactory(commands.Factory{ - Group: "scanner", - Build: func(deps *commands.Deps, reg *commands.CommandRegistry) { - es, _ := deps.EngineSet.(*engine.Set) - if es == nil || es.Gogo == nil { + Capability: "gogo", + Build: func(d *commands.Deps, reg *commands.CommandRegistry) { + es, ok := deps.Get(d.Bag, engine.SetKey) + if !ok || es == nil || es.Gogo == nil { + d.Skip("gogo", deps.Name(engine.SetKey)+".Gogo") return } - impl := New(es.Gogo).WithLogger(deps.GetLogger()).WithProxy(deps.ScannerProxy).WithDataBus(deps.DataBus) - reg.Register(commands.Command{Name: impl.Name(), Usage: impl.Usage(), QuickReference: impl.QuickReference(), Run: impl.Run, SetProxy: impl.SetProxy, GetProxy: func() string { return impl.Proxy }}, "scanner") + impl := New(es.Gogo).WithLogger(d.GetLogger()).WithProxy(d.ScannerProxy).WithDataBus(d.DataBus) + reg.Register(commands.Command{Name: impl.Name(), Usage: impl.Usage(), QuickReference: impl.QuickReference(), Run: impl.Run}, "scanner") }, }) } diff --git a/tools/ioa/commands.go b/tools/ioa/commands.go index 9d939458..e6034a2e 100644 --- a/tools/ioa/commands.go +++ b/tools/ioa/commands.go @@ -9,8 +9,8 @@ import ( "strings" "sync" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" "github.com/chainreactors/ioa/protocols" ) diff --git a/tools/ioa/commands_test.go b/tools/ioa/commands_test.go index b027edb0..f01b33d6 100644 --- a/tools/ioa/commands_test.go +++ b/tools/ioa/commands_test.go @@ -10,7 +10,7 @@ import ( "strings" "testing" - "github.com/chainreactors/aiscan/pkg/agent" + "github.com/chainreactors/aiscan/agent" "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/ioa/protocols" ) diff --git a/tools/ioa/keys.go b/tools/ioa/keys.go new file mode 100644 index 00000000..1fac42d7 --- /dev/null +++ b/tools/ioa/keys.go @@ -0,0 +1,9 @@ +package ioa + +import ( + "github.com/chainreactors/aiscan/core/deps" + "github.com/chainreactors/ioa/protocols" +) + +// ClientKey carries the bound IOA client to the ioa command factory. +var ClientKey = deps.NewKey[protocols.ClientAPI]("ioa.Client") diff --git a/tools/ioa/register.go b/tools/ioa/register.go index f5097b5e..be04f1fc 100644 --- a/tools/ioa/register.go +++ b/tools/ioa/register.go @@ -1,8 +1,9 @@ package ioa import ( + "github.com/chainreactors/aiscan/core/capability" + "github.com/chainreactors/aiscan/core/deps" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/ioa/protocols" _ "github.com/chainreactors/ioa/protocols/checkpoint" _ "github.com/chainreactors/ioa/protocols/handoff" @@ -10,14 +11,19 @@ import ( ) func init() { + capability.Register(capability.Descriptor{ + ID: "ioa", Kind: capability.KindService, Group: "ioa", + Requires: []string{"ioa.ClientAPI"}, + }) commands.RegisterFactory(commands.Factory{ - Group: "ioa", - Build: func(deps *commands.Deps, reg *commands.CommandRegistry) { - client, _ := deps.IOAClient.(protocols.ClientAPI) - if client == nil { + Capability: "ioa", + Build: func(d *commands.Deps, reg *commands.CommandRegistry) { + client, ok := deps.Get(d.Bag, ClientKey) + if !ok || client == nil { + d.Skip("ioa", deps.Name(ClientKey)) return } - for _, cmd := range NewCommands(client, deps.NodeName, deps.NodeMeta) { + for _, cmd := range NewCommands(client, d.NodeName, d.NodeMeta) { reg.Register(cmd, "ioa") } }, diff --git a/tools/katana/katana.go b/tools/katana/katana.go index e52d2e06..a6cd5ce6 100644 --- a/tools/katana/katana.go +++ b/tools/katana/katana.go @@ -13,8 +13,8 @@ import ( "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" "github.com/chainreactors/aiscan/tools/toolargs" "github.com/projectdiscovery/goflags" "github.com/projectdiscovery/gologger" diff --git a/tools/katana/register.go b/tools/katana/register.go index 3d4a08da..404f5a09 100644 --- a/tools/katana/register.go +++ b/tools/katana/register.go @@ -3,12 +3,14 @@ package katana import ( + "github.com/chainreactors/aiscan/core/capability" "github.com/chainreactors/aiscan/pkg/commands" ) func init() { + capability.Register(capability.Descriptor{ID: "katana", Kind: capability.KindScanner, Group: "scanner", CLIName: "katana", Summary: "katana", UsageLine: " katana Run katana web crawler", Usage: func() string { return New().Usage() }, Skills: []string{"katana"}}) commands.RegisterFactory(commands.Factory{ - Group: "scanner", + Capability: "katana", Build: func(deps *commands.Deps, reg *commands.CommandRegistry) { logger := deps.GetLogger() impl := New().WithLogger(logger).WithProxy(deps.ScannerProxy).WithDataBus(deps.DataBus) diff --git a/tools/neutron/neutron.go b/tools/neutron/neutron.go index ca0ac53e..c0078884 100644 --- a/tools/neutron/neutron.go +++ b/tools/neutron/neutron.go @@ -14,8 +14,8 @@ import ( "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" scanengine "github.com/chainreactors/aiscan/tools/scan/engine" "github.com/chainreactors/aiscan/tools/toolargs" "github.com/chainreactors/neutron/templates" diff --git a/tools/neutron/register.go b/tools/neutron/register.go index 9a83bc4a..f2807fc2 100644 --- a/tools/neutron/register.go +++ b/tools/neutron/register.go @@ -1,22 +1,28 @@ package neutron import ( - cfg "github.com/chainreactors/aiscan/core/config" + "github.com/chainreactors/aiscan/core/capability" + "github.com/chainreactors/aiscan/core/deps" "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/aiscan/tools/scan/engine" ) func init() { - cfg.ExtraScannerUsage["neutron"] = func() string { return New(nil, nil).Usage() } + capability.Register(capability.Descriptor{ + ID: "neutron", Kind: capability.KindScanner, Group: "scanner", + CLIName: "neutron", Summary: "neutron", UsageLine: " neutron Run neutron directly", + Usage: func() string { return New(nil, nil).Usage() }, Requires: []string{"scan.engine.Set.Neutron"}, + }) commands.RegisterFactory(commands.Factory{ - Group: "scanner", - Build: func(deps *commands.Deps, reg *commands.CommandRegistry) { - es, _ := deps.EngineSet.(*engine.Set) - if es == nil || es.Neutron == nil { + Capability: "neutron", + Build: func(d *commands.Deps, reg *commands.CommandRegistry) { + es, ok := deps.Get(d.Bag, engine.SetKey) + if !ok || es == nil || es.Neutron == nil { + d.Skip("neutron", deps.Name(engine.SetKey)+".Neutron") return } - impl := New(es.Neutron, es.Index).WithLogger(deps.GetLogger()).WithProxy(deps.ScannerProxy).WithDataBus(deps.DataBus) - reg.Register(commands.Command{Name: impl.Name(), Usage: impl.Usage(), QuickReference: impl.QuickReference(), Run: impl.Run, SetProxy: impl.SetProxy, GetProxy: func() string { return impl.Proxy }}, "scanner") + impl := New(es.Neutron, es.Index).WithLogger(d.GetLogger()).WithProxy(d.ScannerProxy).WithDataBus(d.DataBus) + reg.Register(commands.Command{Name: impl.Name(), Usage: impl.Usage(), QuickReference: impl.QuickReference(), Run: impl.Run}, "scanner") }, }) } diff --git a/tools/neutron/sdk_stage.go b/tools/neutron/sdk_stage.go index e9832d44..7188d41d 100644 --- a/tools/neutron/sdk_stage.go +++ b/tools/neutron/sdk_stage.go @@ -8,7 +8,7 @@ import ( "sync" "time" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/core/telemetry" scanengine "github.com/chainreactors/aiscan/tools/scan/engine" "github.com/chainreactors/neutron/common" "github.com/chainreactors/neutron/templates" diff --git a/tools/passive/passive.go b/tools/passive/passive.go index 6a1f5f73..eda2e4e9 100644 --- a/tools/passive/passive.go +++ b/tools/passive/passive.go @@ -14,8 +14,8 @@ import ( "strings" "time" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" "github.com/chainreactors/aiscan/tools/scan/engine" "github.com/projectdiscovery/uncover/sources" ) diff --git a/tools/passive/register.go b/tools/passive/register.go index e4b6f275..b966821d 100644 --- a/tools/passive/register.go +++ b/tools/passive/register.go @@ -1,28 +1,37 @@ //go:build full +// Passive uses the full-only uncover engine API; the standard stub does not +// expose QueryRaw, RawFofa or RawHunter, so this is a real implementation gate. + package passive import ( - "github.com/chainreactors/aiscan/core/config" + "github.com/chainreactors/aiscan/core/capability" + "github.com/chainreactors/aiscan/core/deps" "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/aiscan/tools/scan/engine" ) func init() { - config.ExtraCommands["passive"] = true - config.ExtraUsageEntries = append(config.ExtraUsageEntries, " passive Run passive cyberspace recon") - config.ExtraSummaryEntries = append(config.ExtraSummaryEntries, "passive") - config.ExtraScannerUsage["passive"] = func() string { return New(nil).Usage() } + capability.Register(capability.Descriptor{ + ID: "passive", Kind: capability.KindScanner, Group: "scanner", + CLIName: "passive", Summary: "passive", UsageLine: " passive Run passive cyberspace recon", + Usage: func() string { return New(nil).Usage() }, Skills: []string{"passive"}, + Requires: []string{"scan.engine.Set.Uncover"}, + }) commands.RegisterFactory(commands.Factory{ - Group: "scanner", - Build: func(deps *commands.Deps, reg *commands.CommandRegistry) { + Capability: "passive", + Build: func(d *commands.Deps, reg *commands.CommandRegistry) { + // passive registers with a nil backend so its usage stays visible; + // every query then reports that no recon source is configured. var backend QueryEngine - if es, ok := deps.EngineSet.(*engine.Set); ok && es != nil && es.Uncover != nil { + if es, ok := deps.Get(d.Bag, engine.SetKey); ok && es != nil && es.Uncover != nil { backend = es.Uncover + } else { + d.Skip("passive.recon", deps.Name(engine.SetKey)+".Uncover") } - logger := deps.GetLogger() - impl := New(backend).WithLogger(logger) + impl := New(backend).WithLogger(d.GetLogger()) reg.Register(commands.Command{Name: impl.Name(), Usage: impl.Usage(), Run: impl.Run}, "scanner") }, }) diff --git a/tools/playwright/advanced.go b/tools/playwright/advanced.go index 47004a09..bbbe0e09 100644 --- a/tools/playwright/advanced.go +++ b/tools/playwright/advanced.go @@ -9,7 +9,7 @@ import ( "strconv" "strings" - "github.com/chainreactors/aiscan/pkg/agent/truncate" + "github.com/chainreactors/aiscan/core/truncate" "github.com/go-rod/rod" "github.com/go-rod/rod/lib/proto" "github.com/ysmood/gson" diff --git a/tools/playwright/browser.go b/tools/playwright/browser.go index 48ba4448..44fd2074 100644 --- a/tools/playwright/browser.go +++ b/tools/playwright/browser.go @@ -14,9 +14,9 @@ import ( "sync" "time" - "github.com/chainreactors/aiscan/pkg/agent/truncate" + "github.com/chainreactors/aiscan/core/telemetry" + "github.com/chainreactors/aiscan/core/truncate" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" "github.com/go-rod/rod" "github.com/go-rod/rod/lib/launcher" "github.com/go-rod/rod/lib/proto" @@ -576,6 +576,12 @@ func (c *Command) getOrLaunchBrowser() (*rod.Browser, error) { Set("disable-dev-shm-usage"). Set("ignore-certificate-errors"). Set("allow-insecure-localhost") + // Prefer an installed browser when one is available. Rod otherwise + // enters its auto-download path, which is inappropriate for offline CI + // and can block even though launcher.LookPath already found Chromium. + if browserPath, ok := launcher.LookPath(); ok { + l = l.Bin(browserPath) + } c.proxyMu.RLock() proxy := c.proxyURL diff --git a/tools/playwright/interact.go b/tools/playwright/interact.go index c807780d..342e9d9b 100644 --- a/tools/playwright/interact.go +++ b/tools/playwright/interact.go @@ -548,7 +548,7 @@ var keyNameMap = map[string]input.Key{ "home": input.Home, "end": input.End, "pageup": input.PageUp, "pagedown": input.PageDown, "insert": input.Insert, - "f1": input.F1, "f2": input.F2, "f3": input.F3, "f4": input.F4, + "f1": input.F1, "f2": input.F2, "f3": input.F3, "f4": input.F4, "f5": input.F5, "f6": input.F6, "f7": input.F7, "f8": input.F8, "f9": input.F9, "f10": input.F10, "f11": input.F11, "f12": input.F12, "shift": input.ShiftLeft, "control": input.ControlLeft, diff --git a/tools/playwright/register.go b/tools/playwright/register.go index 1751e1cf..a1d58326 100644 --- a/tools/playwright/register.go +++ b/tools/playwright/register.go @@ -2,11 +2,15 @@ package playwright -import "github.com/chainreactors/aiscan/pkg/commands" +import ( + "github.com/chainreactors/aiscan/core/capability" + "github.com/chainreactors/aiscan/pkg/commands" +) func init() { + capability.Register(capability.Descriptor{ID: "browser", Kind: capability.KindTool, Group: "browser", Optional: true, Default: true}) commands.RegisterFactory(commands.Factory{ - Group: "browser", + Capability: "browser", Build: func(deps *commands.Deps, reg *commands.CommandRegistry) { impl := New(deps.WorkDir) reg.Register(commands.Command{Name: impl.Name(), Usage: impl.Usage(), Run: impl.Run, Close: impl.Close}, "browser") diff --git a/tools/playwright/session.go b/tools/playwright/session.go index edb25766..fa51ab41 100644 --- a/tools/playwright/session.go +++ b/tools/playwright/session.go @@ -17,8 +17,8 @@ import ( "github.com/go-rod/rod" "github.com/go-rod/rod/lib/proto" "github.com/go-rod/stealth" - "github.com/ysmood/gson" katanajs "github.com/projectdiscovery/katana/pkg/engine/headless/js" + "github.com/ysmood/gson" ) const ( @@ -742,26 +742,26 @@ func (c *Command) execDetach(ctx context.Context, args []string) (string, error) type openOpts struct { commonOpts - sessName string - opTimeout time.Duration - noSpeedUp bool - ignoreHTTPSErrs bool - viewportSize string // "WxH" - geolocation string // "lat,lon" - timezone string - colorScheme string // light|dark - lang string - device string - loadStoragePath string - saveStoragePath string // stored on Session, dumped at close - saveHARPath string // stored on Session, dumped at close - saveHARGlob string - proxyServer string - proxyBypass string - blockSW bool - record bool - headed bool - cdpURL string + sessName string + opTimeout time.Duration + noSpeedUp bool + ignoreHTTPSErrs bool + viewportSize string // "WxH" + geolocation string // "lat,lon" + timezone string + colorScheme string // light|dark + lang string + device string + loadStoragePath string + saveStoragePath string // stored on Session, dumped at close + saveHARPath string // stored on Session, dumped at close + saveHARGlob string + proxyServer string + proxyBypass string + blockSW bool + record bool + headed bool + cdpURL string } func parseOpenOpts(args []string, usage string) (openOpts, error) { @@ -958,13 +958,13 @@ func parseGeolocation(s string) (float64, float64, error) { // storageState mirrors the Playwright storage state format. type storageState struct { - Cookies []json.RawMessage `json:"cookies"` - LocalStorage []localStorageEntry `json:"origins"` + Cookies []json.RawMessage `json:"cookies"` + LocalStorage []localStorageEntry `json:"origins"` } type localStorageEntry struct { - Origin string `json:"origin"` - LocalStorage []nameValuePair `json:"localStorage"` + Origin string `json:"origin"` + LocalStorage []nameValuePair `json:"localStorage"` } type nameValuePair struct { diff --git a/tools/proton/command.go b/tools/proton/command.go index af7adf37..210fbc55 100644 --- a/tools/proton/command.go +++ b/tools/proton/command.go @@ -17,8 +17,8 @@ import ( "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" "github.com/chainreactors/aiscan/tools/toolargs" "github.com/chainreactors/neutron/operators" "github.com/chainreactors/neutron/protocols" diff --git a/tools/proton/register.go b/tools/proton/register.go index bd513802..13783f20 100644 --- a/tools/proton/register.go +++ b/tools/proton/register.go @@ -1,26 +1,30 @@ package proton import ( - cfg "github.com/chainreactors/aiscan/core/config" + "github.com/chainreactors/aiscan/core/capability" + "github.com/chainreactors/aiscan/core/deps" "github.com/chainreactors/aiscan/core/resources" "github.com/chainreactors/aiscan/pkg/commands" ) func init() { - cfg.ExtraCommands["proton"] = true - cfg.ExtraUsageEntries = append(cfg.ExtraUsageEntries, " proton Run proton sensitive info scanner") - cfg.ExtraSummaryEntries = append(cfg.ExtraSummaryEntries, "proton") - cfg.ExtraScannerUsage["proton"] = func() string { return New().Usage() } + capability.Register(capability.Descriptor{ + ID: "proton", Kind: capability.KindScanner, Group: "scanner", + CLIName: "proton", Summary: "proton", UsageLine: " proton Run proton sensitive info scanner", + Usage: func() string { return New().Usage() }, + }) commands.RegisterFactory(commands.Factory{ - Group: "scanner", - Build: func(deps *commands.Deps, reg *commands.CommandRegistry) { - logger := deps.GetLogger() - cmd := New().WithLogger(logger).WithProxy(deps.ScannerProxy).WithDataBus(deps.DataBus) - if rs, ok := deps.Resources.(*resources.Set); ok && rs != nil { + Capability: "proton", + Build: func(d *commands.Deps, reg *commands.CommandRegistry) { + cmd := New().WithLogger(d.GetLogger()).WithProxy(d.ScannerProxy).WithDataBus(d.DataBus) + if rs, ok := deps.Get(d.Bag, resources.SetKey); ok && rs != nil { cmd.WithResourceProvider(rs.ProtonConfig) + } else { + // proton still runs, but only with its built-in rules. + d.Skip("proton.rules", deps.Name(resources.SetKey)) } - cmd.SetWorkDir(deps.WorkDir) - reg.Register(commands.Command{Name: cmd.Name(), Usage: cmd.Usage(), Run: cmd.Run, SetProxy: cmd.SetProxy, GetProxy: func() string { return cmd.Proxy }}, "scanner") + cmd.SetWorkDir(d.WorkDir) + reg.Register(commands.Command{Name: cmd.Name(), Usage: cmd.Usage(), Run: cmd.Run}, "scanner") }, }) } diff --git a/tools/proton/register_test.go b/tools/proton/register_test.go index 78e8a98b..b303847e 100644 --- a/tools/proton/register_test.go +++ b/tools/proton/register_test.go @@ -4,12 +4,13 @@ import ( "slices" "testing" + "github.com/chainreactors/aiscan/core/capability" "github.com/chainreactors/aiscan/pkg/commands" ) func TestFactoryBuildsProtonWithScannerGroup(t *testing.T) { registry := commands.NewRegistry() - commands.BuildGroup("scanner", &commands.Deps{WorkDir: t.TempDir()}, registry) + commands.BuildPlan(capability.Select(capability.Options{Groups: []string{"scanner"}}), &commands.Deps{WorkDir: t.TempDir()}, registry) if !registry.Has("proton") { t.Fatal("scanner group did not register proton") diff --git a/tools/proxy/command.go b/tools/proxy/command.go index fd1634fa..88ddfdb9 100644 --- a/tools/proxy/command.go +++ b/tools/proxy/command.go @@ -6,8 +6,8 @@ import ( "net/url" "strings" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" "github.com/chainreactors/proxyclient" "github.com/chainreactors/proxyclient/extra/clash" goflags "github.com/jessevdk/go-flags" diff --git a/tools/proxy/mitm.go b/tools/proxy/mitm.go index 1ecbf8cc..6e28703d 100644 --- a/tools/proxy/mitm.go +++ b/tools/proxy/mitm.go @@ -9,8 +9,8 @@ import ( "sync" "time" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" mitmproxy "github.com/chainreactors/utils/mitmproxy/proxy" goflags "github.com/jessevdk/go-flags" ) diff --git a/tools/proxy/mitm_test.go b/tools/proxy/mitm_test.go index f737e5e3..f9dd2631 100644 --- a/tools/proxy/mitm_test.go +++ b/tools/proxy/mitm_test.go @@ -177,9 +177,6 @@ func TestMITMCapture_NonHTTP_Fallback(t *testing.T) { } func TestMITMCapture_ServerFirst_Fallback(t *testing.T) { - if raceEnabled { - t.Skip("flaky under -race: mitmproxy internal goroutine scheduling causes i/o timeout on CI") - } // Server-first protocol (like SSH): server sends banner, client waits. // MITM should timeout on Peek and fallback to raw transfer. tcpServer, err := net.Listen("tcp", "127.0.0.1:0") @@ -187,16 +184,19 @@ func TestMITMCapture_ServerFirst_Fallback(t *testing.T) { t.Fatal(err) } defer tcpServer.Close() + serverReady := make(chan error, 1) go func() { - for { - conn, err := tcpServer.Accept() - if err != nil { - return - } - conn.Write([]byte("SSH-2.0-TestServer\r\n")) + conn, err := tcpServer.Accept() + if err != nil { + serverReady <- err + return + } + defer conn.Close() + _, err = conn.Write([]byte("SSH-2.0-TestServer\r\n")) + serverReady <- err + if err == nil { buf := make([]byte, 256) - conn.Read(buf) - conn.Close() + _, _ = conn.Read(buf) } }() @@ -209,31 +209,49 @@ func TestMITMCapture_ServerFirst_Fallback(t *testing.T) { t.Fatal(err) } - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() conn, err := dial(ctx, "tcp", tcpServer.Addr().String()) if err != nil { t.Fatal(err) } defer conn.Close() + select { + case err := <-serverReady: + if err != nil { + t.Fatalf("server banner write: %v", err) + } + case <-ctx.Done(): + t.Fatalf("server did not accept proxy connection: %v", ctx.Err()) + } buf := make([]byte, 64) - // The banner only arrives after the MITM's ~3s peek timeout expires and it - // falls back to raw transfer. Keep this deadline well above that timeout so - // scheduling jitter under -race / CI load can't race it (was 8s → flaky). - conn.SetReadDeadline(time.Now().Add(30 * time.Second)) - n, err := io.ReadAtLeast(conn, buf, 3) - if err != nil { - t.Fatalf("expected SSH banner data, got error: %v", err) + type readResult struct { + n int + err error + } + readDone := make(chan readResult, 1) + go func() { + n, err := io.ReadAtLeast(conn, buf, 3) + readDone <- readResult{n: n, err: err} + }() + var read readResult + select { + case read = <-readDone: + case <-ctx.Done(): + t.Fatalf("proxy did not enter raw fallback after synchronized banner write: %v", ctx.Err()) + } + if read.err != nil { + t.Fatalf("expected SSH banner data, got error: %v", read.err) } - banner := string(buf[:n]) + banner := string(buf[:read.n]) if !strings.Contains(banner, "SSH-") && !strings.Contains(banner, "SH-") { t.Fatalf("expected SSH banner fragment, got %q", banner) } if store.Count() != 0 { t.Fatalf("server-first protocol should not capture flows, got %d", store.Count()) } - t.Logf("server-first fallback OK: received %q, 0 flows captured", string(buf[:n])) + t.Logf("server-first fallback OK: received %q, 0 flows captured", string(buf[:read.n])) } // === Latency Benchmark === @@ -387,7 +405,7 @@ func BenchmarkFlowStore_Query(b *testing.B) { store.Add(Flow{ Method: "GET", URL: fmt.Sprintf("http://host%d.com/path%d", i%10, i), - StatusCode: 200 + (i % 5) * 100, + StatusCode: 200 + (i%5)*100, Host: fmt.Sprintf("host%d.com", i%10), }) } diff --git a/tools/proxy/race_norace_test.go b/tools/proxy/race_norace_test.go deleted file mode 100644 index 84a78f40..00000000 --- a/tools/proxy/race_norace_test.go +++ /dev/null @@ -1,5 +0,0 @@ -//go:build !race - -package proxy - -const raceEnabled = false diff --git a/tools/proxy/race_test.go b/tools/proxy/race_test.go deleted file mode 100644 index b624eb94..00000000 --- a/tools/proxy/race_test.go +++ /dev/null @@ -1,5 +0,0 @@ -//go:build race - -package proxy - -const raceEnabled = true diff --git a/tools/proxy/register_command.go b/tools/proxy/register_command.go index 1a238201..43e9ebf4 100644 --- a/tools/proxy/register_command.go +++ b/tools/proxy/register_command.go @@ -6,6 +6,7 @@ import ( "net/url" "strings" + "github.com/chainreactors/aiscan/core/capability" "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/proxyclient" @@ -18,8 +19,9 @@ import ( ) func init() { + capability.Register(capability.Descriptor{ID: "proxy", Kind: capability.KindService, Group: "proxy"}) commands.RegisterFactory(commands.Factory{ - Group: "proxy", + Capability: "proxy", Build: func(deps *commands.Deps, reg *commands.CommandRegistry) { state := NewState(deps.ScannerProxy) cmd := New(state) diff --git a/tools/proxy/state.go b/tools/proxy/state.go index 6e10c17f..43270819 100644 --- a/tools/proxy/state.go +++ b/tools/proxy/state.go @@ -23,8 +23,8 @@ type State struct { subscribeURL string activeNode *clash.ProxyNode activeURL string - autoURL string // clash:// URL for auto mode - autoDial proxyclient.Dial // pre-built dial for auto mode + autoURL string // clash:// URL for auto mode + autoDial proxyclient.Dial // pre-built dial for auto mode } func NewState(originalProxy string) *State { @@ -151,7 +151,7 @@ func (s *State) TestNode(ctx context.Context, node *clash.ProxyNode) (time.Durat DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) { return dial.DialContext(ctx, network, addr) }, - TLSClientConfig: &tls.Config{}, + TLSClientConfig: &tls.Config{}, DisableKeepAlives: true, } client := &http.Client{ diff --git a/tools/register_command.go b/tools/register_command.go index 33bf7f10..a75859d4 100644 --- a/tools/register_command.go +++ b/tools/register_command.go @@ -1,39 +1,40 @@ package tools import ( - cfg "github.com/chainreactors/aiscan/core/config" + "github.com/chainreactors/aiscan/core/capability" + "github.com/chainreactors/aiscan/core/deps" "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/aiscan/tools/scan" "github.com/chainreactors/aiscan/tools/scan/engine" ) func init() { - cfg.ExtraScannerUsage["scan"] = scan.Usage + capability.Register(capability.Descriptor{ + ID: "scan", Kind: capability.KindScanner, Group: "scanner", + CLIName: "scan", Summary: "scan", Usage: scan.Usage, + Requires: []string{"scan.engine.Set.Gogo", "scan.engine.Set.Spray"}, + }) commands.RegisterFactory(commands.Factory{ - Group: "scanner", - Build: func(deps *commands.Deps, reg *commands.CommandRegistry) { - es, _ := deps.EngineSet.(*engine.Set) - if es == nil { + Capability: "scan", + Build: func(d *commands.Deps, reg *commands.CommandRegistry) { + es, ok := deps.Get(d.Bag, engine.SetKey) + if !ok || es == nil || es.Gogo == nil || es.Spray == nil { + d.Skip("scan", deps.Name(engine.SetKey)+".Gogo+.Spray") return } - var scanOpts []scan.Option - for _, o := range deps.ScanOpts { - if opt, ok := o.(scan.Option); ok { - scanOpts = append(scanOpts, opt) - } - } - if deps.ScannerProxy != "" { - scanOpts = append(scanOpts, scan.WithProxy(deps.ScannerProxy)) + // copy: the bag's slice is shared with every other build + stored, _ := deps.Get(d.Bag, scan.OptsKey) + scanOpts := append([]scan.Option(nil), stored...) + if d.ScannerProxy != "" { + scanOpts = append(scanOpts, scan.WithProxy(d.ScannerProxy)) } - if deps.DataBus != nil { - scanOpts = append(scanOpts, scan.WithDataBus(deps.DataBus)) + if d.DataBus != nil { + scanOpts = append(scanOpts, scan.WithDataBus(d.DataBus)) } - if es.Gogo != nil && es.Spray != nil { - impl := scan.New(es, scanOpts...) - reg.Register(commands.Command{Name: impl.Name(), Usage: impl.Usage(), Run: impl.Run}, "scanner") - } + impl := scan.New(es, scanOpts...) + reg.Register(commands.Command{Name: impl.Name(), Usage: impl.Usage(), Run: impl.Run}, "scanner") }, }) } diff --git a/tools/register_command_integration_test.go b/tools/register_command_integration_test.go index 4f31e33e..0f5478de 100644 --- a/tools/register_command_integration_test.go +++ b/tools/register_command_integration_test.go @@ -13,8 +13,8 @@ import ( "testing" "time" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" passivecmd "github.com/chainreactors/aiscan/tools/passive" "github.com/chainreactors/aiscan/tools/scan/engine" ) diff --git a/tools/register_command_test.go b/tools/register_command_test.go index 5c6ad7d4..5d7ffd5c 100644 --- a/tools/register_command_test.go +++ b/tools/register_command_test.go @@ -29,11 +29,10 @@ import ( func buildRegistry(engineSet *engine.Set) *commands.CommandRegistry { reg := commands.NewRegistry() - deps := &commands.Deps{ - EngineSet: engineSet, - Resources: engineSet.Resources, - } - commands.BuildAll(deps, reg) + deps := &commands.Deps{} + commands.Provide(deps, engine.SetKey, engineSet) + commands.Provide(deps, resources.SetKey, engineSet.Resources) + buildTestGroups([]string{"scanner", "search"}, deps, reg) return reg } diff --git a/tools/scan/collector.go b/tools/scan/collector.go index f8854e9b..c0a81569 100644 --- a/tools/scan/collector.go +++ b/tools/scan/collector.go @@ -168,7 +168,14 @@ func (c *collector) TerminalString(color bool) string { } func (c *collector) ReportMarkdown() string { - return formatMarkdown(c) + return output.RenderReport(c.StructuredResult(), output.ReportOptions{ + Style: output.StyleMarkdown, + Title: "Scan Report", + Sitemap: true, + CollapseBare: true, + Metrics: true, + Inventory: true, + }) } func (c *collector) JSONLines() (string, error) { diff --git a/tools/scan/command.go b/tools/scan/command.go index c405df38..de4c5800 100644 --- a/tools/scan/command.go +++ b/tools/scan/command.go @@ -7,12 +7,12 @@ import ( "os" "path/filepath" + "github.com/chainreactors/aiscan/agent" + "github.com/chainreactors/aiscan/core/aop" "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" - "github.com/chainreactors/aiscan/pkg/agent" - "github.com/chainreactors/aiscan/pkg/aop" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" "github.com/chainreactors/aiscan/tools/scan/engine" "github.com/chainreactors/aiscan/tools/scan/pipeline" "github.com/chainreactors/aiscan/tools/toolargs" diff --git a/tools/scan/command_test.go b/tools/scan/command_test.go index 5bf8967a..5b91057c 100644 --- a/tools/scan/command_test.go +++ b/tools/scan/command_test.go @@ -16,8 +16,8 @@ import ( "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" "github.com/chainreactors/aiscan/tools/scan/engine" "github.com/chainreactors/aiscan/tools/scan/pipeline" "github.com/chainreactors/fingers/common" diff --git a/tools/scan/engine/gogo.go b/tools/scan/engine/gogo.go index c5e9460c..9c097a7e 100644 --- a/tools/scan/engine/gogo.go +++ b/tools/scan/engine/gogo.go @@ -5,11 +5,11 @@ import ( "fmt" "os" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/core/telemetry" gogopkg "github.com/chainreactors/gogo/v2/pkg" - "github.com/chainreactors/utils/parsers" "github.com/chainreactors/sdk/gogo" sdktypes "github.com/chainreactors/sdk/pkg/types" + "github.com/chainreactors/utils/parsers" ) const GogoTempLogFile = ".sock.lock" diff --git a/tools/scan/engine/keys.go b/tools/scan/engine/keys.go new file mode 100644 index 00000000..6356e063 --- /dev/null +++ b/tools/scan/engine/keys.go @@ -0,0 +1,7 @@ +package engine + +import "github.com/chainreactors/aiscan/core/deps" + +// SetKey carries the initialized scanner engines to the command factories. +// Declared here so pkg/commands never has to link the scanner SDKs. +var SetKey = deps.NewKey[*Set]("scan.engine.Set") diff --git a/tools/scan/engine/neutron.go b/tools/scan/engine/neutron.go index cf9bffe2..ff2a43df 100644 --- a/tools/scan/engine/neutron.go +++ b/tools/scan/engine/neutron.go @@ -5,7 +5,7 @@ import ( "errors" "fmt" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/neutron/common" "github.com/chainreactors/neutron/templates" "github.com/chainreactors/sdk/neutron" diff --git a/tools/scan/engine/race_norace_test.go b/tools/scan/engine/race_norace_test.go deleted file mode 100644 index 9f576e43..00000000 --- a/tools/scan/engine/race_norace_test.go +++ /dev/null @@ -1,5 +0,0 @@ -//go:build !race - -package engine - -const raceEnabled = false diff --git a/tools/scan/engine/race_test.go b/tools/scan/engine/race_test.go deleted file mode 100644 index ee11e288..00000000 --- a/tools/scan/engine/race_test.go +++ /dev/null @@ -1,5 +0,0 @@ -//go:build race - -package engine - -const raceEnabled = true diff --git a/tools/scan/engine/set.go b/tools/scan/engine/set.go index 61e744f8..66587bcf 100644 --- a/tools/scan/engine/set.go +++ b/tools/scan/engine/set.go @@ -9,8 +9,8 @@ import ( "sync" "github.com/chainreactors/aiscan/core/resources" - "github.com/chainreactors/aiscan/pkg/telemetry" - "github.com/chainreactors/aiscan/pkg/util" + "github.com/chainreactors/aiscan/core/telemetry" + "github.com/chainreactors/aiscan/core/util" "github.com/chainreactors/fingers/alias" fingersLib "github.com/chainreactors/fingers/fingers" neutronhttp "github.com/chainreactors/neutron/protocols/http" diff --git a/tools/scan/engine/set_test.go b/tools/scan/engine/set_test.go index 2b2f42ed..682d0362 100644 --- a/tools/scan/engine/set_test.go +++ b/tools/scan/engine/set_test.go @@ -6,7 +6,7 @@ import ( "testing" "time" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/core/telemetry" gogopkg "github.com/chainreactors/gogo/v2/pkg" "github.com/chainreactors/neutron/operators" neutronhttp "github.com/chainreactors/neutron/protocols/http" @@ -468,9 +468,6 @@ func TestSprayStatsHandlerSafeAfterCancel(t *testing.T) { } func TestZombieStatsHandlerSafeAfterCancel(t *testing.T) { - if raceEnabled { - t.Skip("zombie engine has known races under -race detector") - } eng, err := sdkzombie.NewEngine(nil) if err != nil { t.Fatalf("NewEngine: %v", err) diff --git a/tools/scan/engine/set_uncover_recon.go b/tools/scan/engine/set_uncover_recon.go index 778f9072..530f9c89 100644 --- a/tools/scan/engine/set_uncover_recon.go +++ b/tools/scan/engine/set_uncover_recon.go @@ -5,7 +5,7 @@ package engine import ( "strings" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/core/telemetry" ) func (e *Set) SetupUncover(opts ReconOptions, logger telemetry.Logger) { diff --git a/tools/scan/engine/set_uncover_stub.go b/tools/scan/engine/set_uncover_stub.go index 8e187b89..e1a57bfc 100644 --- a/tools/scan/engine/set_uncover_stub.go +++ b/tools/scan/engine/set_uncover_stub.go @@ -2,6 +2,6 @@ package engine -import "github.com/chainreactors/aiscan/pkg/telemetry" +import "github.com/chainreactors/aiscan/core/telemetry" func (e *Set) SetupUncover(_ ReconOptions, _ telemetry.Logger) {} diff --git a/tools/scan/engine/spray.go b/tools/scan/engine/spray.go index 662be303..b340266c 100644 --- a/tools/scan/engine/spray.go +++ b/tools/scan/engine/spray.go @@ -3,14 +3,21 @@ package engine import ( "context" "fmt" + "sync" "time" - "github.com/chainreactors/aiscan/pkg/telemetry" - "github.com/chainreactors/utils/parsers" + "github.com/chainreactors/aiscan/core/telemetry" sdktypes "github.com/chainreactors/sdk/pkg/types" "github.com/chainreactors/sdk/spray" + "github.com/chainreactors/utils/parsers" ) +// spray's runner construction mutates shared logger/option state inside the +// upstream engine. A scan may schedule check, crawl and plugin capabilities in +// parallel against the same engine, so keep one invocation active at a time +// until its result stream is fully drained. +var sprayExecutionMu sync.Mutex + type SprayCheckOptions struct { URLs []string Host string @@ -40,6 +47,7 @@ func SprayCheckStream(ctx context.Context, eng *spray.Engine, opts SprayCheckOpt if eng == nil { return nil, fmt.Errorf("spray engine is not available") } + sprayExecutionMu.Lock() if opts.Debug { telemetry.EnableLogsDebug() } @@ -58,12 +66,14 @@ func SprayCheckStream(ctx context.Context, eng *spray.Engine, opts SprayCheckOpt } if err != nil { cancel() + sprayExecutionMu.Unlock() return nil, err } out := make(chan *parsers.SprayResult) go func() { defer telemetry.SDKGoRecover("spray") + defer sprayExecutionMu.Unlock() defer cancel() defer close(out) for { diff --git a/tools/scan/engine/uncover.go b/tools/scan/engine/uncover.go index 73e6a517..f8b85b63 100644 --- a/tools/scan/engine/uncover.go +++ b/tools/scan/engine/uncover.go @@ -13,7 +13,7 @@ import ( "strings" "time" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/projectdiscovery/uncover/sources" ) diff --git a/tools/scan/engine/zombie.go b/tools/scan/engine/zombie.go index 3c09bb98..f651de67 100644 --- a/tools/scan/engine/zombie.go +++ b/tools/scan/engine/zombie.go @@ -4,10 +4,10 @@ import ( "context" "fmt" - "github.com/chainreactors/aiscan/pkg/telemetry" - "github.com/chainreactors/utils/parsers" + "github.com/chainreactors/aiscan/core/telemetry" sdktypes "github.com/chainreactors/sdk/pkg/types" sdkzombie "github.com/chainreactors/sdk/zombie" + "github.com/chainreactors/utils/parsers" ) type ZombieWeakpassOptions struct { diff --git a/tools/scan/event.go b/tools/scan/event.go index fc4db17d..b35a1c09 100644 --- a/tools/scan/event.go +++ b/tools/scan/event.go @@ -7,8 +7,8 @@ import ( "sync/atomic" "github.com/chainreactors/aiscan/core/output" - "github.com/chainreactors/utils/parsers" sdktypes "github.com/chainreactors/sdk/pkg/types" + "github.com/chainreactors/utils/parsers" ) type eventKind string diff --git a/tools/scan/http_auth.go b/tools/scan/http_auth.go index 05b062fe..567ca7b3 100644 --- a/tools/scan/http_auth.go +++ b/tools/scan/http_auth.go @@ -62,7 +62,7 @@ func httpAuthClient(timeoutSeconds int) *http.Client { if timeoutSeconds <= 0 { timeoutSeconds = 5 } - transport := http.DefaultTransport.(*http.Transport).Clone() //nolint:errcheck // DefaultTransport is always *http.Transport + transport := http.DefaultTransport.(*http.Transport).Clone() //nolint:errcheck // DefaultTransport is always *http.Transport transport.TLSClientConfig = &tls.Config{InsecureSkipVerify: true} //nolint:gosec // scanner probes must tolerate self-signed certs return &http.Client{ Timeout: time.Duration(timeoutSeconds) * time.Second, diff --git a/tools/scan/input.go b/tools/scan/input.go index 3bdfe620..16a20220 100644 --- a/tools/scan/input.go +++ b/tools/scan/input.go @@ -8,9 +8,9 @@ import ( "os" "strings" - "github.com/chainreactors/utils/parsers" sdkzombie "github.com/chainreactors/sdk/zombie" "github.com/chainreactors/utils" + "github.com/chainreactors/utils/parsers" zombiepkg "github.com/chainreactors/zombie/pkg" ) diff --git a/tools/scan/jsonl_writer.go b/tools/scan/jsonl_writer.go index b634d555..e56d2d17 100644 --- a/tools/scan/jsonl_writer.go +++ b/tools/scan/jsonl_writer.go @@ -4,9 +4,9 @@ import ( "encoding/json" "strings" + "github.com/chainreactors/aiscan/core/aop" "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" - "github.com/chainreactors/aiscan/pkg/aop" "github.com/chainreactors/aiscan/tools/scan/pipeline" ) diff --git a/tools/scan/keys.go b/tools/scan/keys.go new file mode 100644 index 00000000..2ceadd8f --- /dev/null +++ b/tools/scan/keys.go @@ -0,0 +1,7 @@ +package scan + +import "github.com/chainreactors/aiscan/core/deps" + +// OptsKey carries scan options built by the assembly layer (parent agent, deep +// browser, skill reader) that cannot be expressed as plain Deps fields. +var OptsKey = deps.NewKey[[]Option]("scan.Options") diff --git a/tools/scan/options.go b/tools/scan/options.go index 99f7d36d..cc82d83d 100644 --- a/tools/scan/options.go +++ b/tools/scan/options.go @@ -3,10 +3,10 @@ package scan import ( "context" + "github.com/chainreactors/aiscan/agent" "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" - "github.com/chainreactors/aiscan/pkg/agent" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/core/telemetry" ) type Option func(*Command) diff --git a/tools/scan/pipeline/pipeline.go b/tools/scan/pipeline/pipeline.go index 17b63de9..b2999122 100644 --- a/tools/scan/pipeline/pipeline.go +++ b/tools/scan/pipeline/pipeline.go @@ -6,7 +6,7 @@ import ( "sync" "github.com/chainreactors/aiscan/core/eventbus" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/core/telemetry" ) type Event interface { diff --git a/tools/scan/report.go b/tools/scan/report.go index aeb4e018..035645ed 100644 --- a/tools/scan/report.go +++ b/tools/scan/report.go @@ -1,15 +1,12 @@ package scan import ( - "fmt" - "sort" "strconv" "strings" "time" "github.com/chainreactors/aiscan/core/output" "github.com/chainreactors/utils/parsers" - sdktypes "github.com/chainreactors/sdk/pkg/types" ) func formatSummary(d *collector, color bool) string { @@ -25,104 +22,10 @@ func formatSummary(d *collector, color bool) string { } } sb.WriteString(formatScanSummaryLine(d, stats, color)) - - if len(d.trace) > 0 { - for _, line := range d.trace { - sb.WriteString(line) - sb.WriteString("\n") - } - } - - return sb.String() -} - -func formatMarkdown(d *collector) string { - d.mu.Lock() - defer d.mu.Unlock() - stats := d.statsSnapshotLocked() - - var sb strings.Builder - sb.WriteString("# Scan Report\n\n") - sb.WriteString(formatScanSummaryLine(d, stats, false)) - sb.WriteString("\n\n") - - sb.WriteString("## Metrics\n\n") - sb.WriteString("| Metric | Value |\n") - sb.WriteString("| --- | ---: |\n") - sb.WriteString(fmt.Sprintf("| Inputs | %d |\n", stats.Inputs)) - sb.WriteString(fmt.Sprintf("| Open services | %d |\n", len(d.gogoResults))) - sb.WriteString(fmt.Sprintf("| Web endpoints | %d |\n", len(d.seenWeb))) - sb.WriteString(fmt.Sprintf("| Web probes | %d |\n", len(d.sprayResults))) - sb.WriteString(fmt.Sprintf("| Fingerprints | %d |\n", len(d.seenFinger))) - sb.WriteString(fmt.Sprintf("| Loots | %d |\n", len(d.loots))) - sb.WriteString(fmt.Sprintf("| Errors | %d |\n", len(d.errors))) - sb.WriteString(fmt.Sprintf("| Tasks | %d |\n", stats.Tasks)) - sb.WriteString(fmt.Sprintf("| Requests | %d |\n", stats.Requests)) - sb.WriteString(fmt.Sprintf("| Duration | %s |\n", stats.Duration().Round(time.Millisecond))) - - if d.debug && len(stats.CapabilityRuns) > 0 { - sb.WriteString("\n## Capability Runs\n\n") - writeCountTable(&sb, "Capability", stats.CapabilityRuns) - } - - if d.debug && len(stats.EngineStats) > 0 { - sb.WriteString("\n## Engine Stats\n\n") - writeEngineStatsTable(&sb, stats.EngineStats) - } - - if len(d.gogoResults) > 0 { - sb.WriteString("\n## Open Services\n\n") - for _, result := range sortedCopy(d.gogoResults, func(a, b *parsers.GOGOResult) bool { - return a.GetTarget() < b.GetTarget() - }) { - writeMarkdownEventLine(&sb, targetEvent(capGogoPortscan, "", newServiceTarget("", result))) - } - } - - if len(d.sprayResults) > 0 { - sb.WriteString("\n## Web Evidence\n\n") - for _, item := range sortedCopy(d.sprayResults, func(a, b sprayObservation) bool { - return sprayResultSortKey(a) < sprayResultSortKey(b) - }) { - if item.Result == nil { - continue - } - writeMarkdownEventLine(&sb, targetEvent(item.Capability, "", newWebProbeTarget("", item.Capability, "", item.Result))) - } - } - - if len(d.loots) > 0 { - sb.WriteString("\n## Loots\n\n") - for _, loot := range sortedCopy(d.loots, func(a, b output.Loot) bool { - if a.Kind != b.Kind { - return a.Kind < b.Kind - } - return a.Description < b.Description - }) { - status, _ := loot.Data["verification_status"].(string) - line := formatEventLine(lootEvent(loot.Kind, loot), false) - if line != "" { - writeMarkdownStatusLine(&sb, line, status) - } - } - } - - if len(d.errors) > 0 { - sb.WriteString("\n## Errors\n\n") - for _, line := range sortedCopy(d.errors, func(a, b string) bool { return a < b }) { - writeMarkdownEventLine(&sb, errorEventOf("scan", line)) - } - } - - if d.debug && len(d.trace) > 0 { - sb.WriteString("\n## Trace\n\n") - for _, line := range d.trace { - sb.WriteString("- ") - sb.WriteString(line) - sb.WriteString("\n") - } + for _, line := range d.trace { + sb.WriteString(line) + sb.WriteString("\n") } - return sb.String() } @@ -139,8 +42,7 @@ func formatScanSummaryLine(d *collector, stats statsSnapshot, color bool) string parts = appendCount64(parts, stats.Requests, "request", "requests") parts = append(parts, stats.Duration().Round(time.Millisecond).String()) c := output.NewColor(color) - body := strings.Join(parts, " ") - return output.FormatLine(output.OutputPrefix("summary", c.Dim), body, c) + "\n" + return output.FormatLine(output.OutputPrefix("summary", c.Dim), strings.Join(parts, " "), c) + "\n" } func appendCount(parts []string, n int, singular, plural string) []string { @@ -159,19 +61,6 @@ func appendCount64(parts []string, n int64, singular, plural string) []string { return append(parts, strconv.FormatInt(n, 10), word) } -func sortedCopy[T any](items []T, less func(a, b T) bool) []T { - out := append([]T(nil), items...) - sort.SliceStable(out, func(i, j int) bool { return less(out[i], out[j]) }) - return out -} - -func sprayResultSortKey(item sprayObservation) string { - if item.Result == nil { - return item.Capability - } - return item.Result.UrlString + "|" + item.Capability + "|" + item.Result.Source.Name() -} - func formatTraceEvent(event pipelineEvent) string { parts := []string{string(event.Action)} if event.Capability != "" { @@ -185,27 +74,21 @@ func formatTraceEvent(event pipelineEvent) string { hostHeader := "" switch target := event.Event.Target.(type) { case scanTarget: - if target.Target != "" { - targetValue = target.Target - } + targetValue = target.Target case serviceTarget: if target.Result != nil { targetValue = target.Result.GetTarget() } case webTarget: - if target.URL != "" { - targetValue = target.URL - } + targetValue = target.URL hostHeader = target.HostHeader case webProbeTarget: - if target.Result != nil && target.Result.UrlString != "" { + if target.Result != nil { targetValue = target.Result.UrlString } hostHeader = target.HostHeader case pocTarget: - if target.Target != "" { - targetValue = target.Target - } + targetValue = target.Target case weakpassTarget: if target.Target.Address() != ":" { targetValue = target.Target.Address() @@ -222,86 +105,3 @@ func formatTraceEvent(event pipelineEvent) string { } return output.FormatLine("[trace]", parsers.JoinOutput(parts...), output.NewColor(false)) } - -func writeMarkdownEventLine(sb *strings.Builder, event event) { - line := formatEventLine(event, false) - if line == "" { - return - } - writeMarkdownStatusLine(sb, line, "") -} - -func writeMarkdownStatusLine(sb *strings.Builder, line, status string) { - if line == "" { - return - } - sb.WriteString("- ") - switch status { - case "not_confirmed": - sb.WriteString("~~") - sb.WriteString(line) - sb.WriteString("~~ *(not confirmed)*") - case "confirmed": - sb.WriteString("**[verified]** ") - sb.WriteString(line) - case "inconclusive": - sb.WriteString("**[inconclusive]** ") - sb.WriteString(line) - case "failed": - sb.WriteString("**[verification failed]** ") - sb.WriteString(line) - default: - sb.WriteString(line) - } - sb.WriteString("\n") -} - - -func sortedMapKeys(values map[string]int) []string { - keys := make([]string, 0, len(values)) - for key := range values { - if key != "" { - keys = append(keys, key) - } - } - sort.Strings(keys) - return keys -} - -func writeCountTable(sb *strings.Builder, label string, values map[string]int) { - sb.WriteString(fmt.Sprintf("| %s | Count |\n", label)) - sb.WriteString("| --- | ---: |\n") - for _, key := range sortedMapKeys(values) { - sb.WriteString(fmt.Sprintf("| %s | %d |\n", key, values[key])) - } -} - -func sortedStatsKeys(values map[string]sdktypes.Stats) []string { - keys := make([]string, 0, len(values)) - for key := range values { - if key != "" { - keys = append(keys, key) - } - } - sort.Strings(keys) - return keys -} - -func writeEngineStatsTable(sb *strings.Builder, values map[string]sdktypes.Stats) { - sb.WriteString("| Source | Engine | Task | Targets | Tasks | Requests | Results | Errors | Duration |\n") - sb.WriteString("| --- | --- | --- | ---: | ---: | ---: | ---: | ---: | ---: |\n") - for _, key := range sortedStatsKeys(values) { - stats := values[key] - sb.WriteString(fmt.Sprintf("| %s | %s | %s | %d | %d | %d | %d | %d | %s |\n", - key, - stats.Engine, - stats.Task, - stats.Targets, - stats.Tasks, - stats.Requests, - stats.Results, - stats.Errors, - stats.Duration.Round(time.Millisecond), - )) - } -} diff --git a/tools/scan/target.go b/tools/scan/target.go index fbc510c2..7ac3a1f7 100644 --- a/tools/scan/target.go +++ b/tools/scan/target.go @@ -5,9 +5,9 @@ import ( "sort" "strings" - "github.com/chainreactors/utils/parsers" sdkzombie "github.com/chainreactors/sdk/zombie" "github.com/chainreactors/utils" + "github.com/chainreactors/utils/parsers" ) type target interface { diff --git a/tools/scan/verify.go b/tools/scan/verify.go index 3fa0c096..c36c3519 100644 --- a/tools/scan/verify.go +++ b/tools/scan/verify.go @@ -5,9 +5,9 @@ import ( "fmt" "strings" + "github.com/chainreactors/aiscan/agent" "github.com/chainreactors/aiscan/core/output" - "github.com/chainreactors/aiscan/pkg/agent" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/core/telemetry" ) type indexedLoot struct { diff --git a/tools/search/fetch.go b/tools/search/fetch.go index a36dc857..f5bdd777 100644 --- a/tools/search/fetch.go +++ b/tools/search/fetch.go @@ -12,9 +12,9 @@ import ( "time" "unicode" - "github.com/chainreactors/aiscan/pkg/agent/truncate" + "github.com/chainreactors/aiscan/core/telemetry" + "github.com/chainreactors/aiscan/core/truncate" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" ) const ( diff --git a/tools/search/register.go b/tools/search/register.go index 5d9a4817..90804a9d 100644 --- a/tools/search/register.go +++ b/tools/search/register.go @@ -1,42 +1,46 @@ package search import ( + "github.com/chainreactors/aiscan/core/capability" + "github.com/chainreactors/aiscan/core/deps" "github.com/chainreactors/aiscan/core/resources" - "github.com/chainreactors/aiscan/pkg/agent/provider" "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/aiscan/tools/scan/engine" "github.com/chainreactors/sdk/pkg/association" ) func init() { + capability.Register(capability.Descriptor{ + ID: "search", Kind: capability.KindTool, Group: "search", + Optional: true, Default: true, + }) commands.RegisterFactory(commands.Factory{ - Group: "search", - Build: func(deps *commands.Deps, reg *commands.CommandRegistry) { - var p provider.Provider - if deps.Provider != nil { - p, _ = deps.Provider.(provider.Provider) - } - - tavily := NewTavilySearch(deps.TavilyKeys) - if deps.ScannerProxy != "" { - tavily.SetProxy(deps.ScannerProxy) + Capability: "search", + Build: func(d *commands.Deps, reg *commands.CommandRegistry) { + tavily := NewTavilySearch(d.TavilyKeys) + if d.ScannerProxy != "" { + tavily.SetProxy(d.ScannerProxy) } - reg.RegisterTool(NewWebSearchTool(p, tavily)) + reg.RegisterTool(NewWebSearchTool(d.Provider, tavily)) fetch := NewFetchCommand() reg.Register(commands.Command{Name: fetch.Name(), Usage: fetch.Usage(), Run: fetch.Run}, "search") var idx *association.Index - if es, ok := deps.EngineSet.(*engine.Set); ok && es != nil { + if es, ok := deps.Get(d.Bag, engine.SetKey); ok && es != nil { idx = es.Index } if idx == nil { - if rs, ok := deps.Resources.(*resources.Set); ok && rs != nil && rs.FingersConfig != nil { + if rs, ok := deps.Get(d.Bag, resources.SetKey); ok && rs != nil && rs.FingersConfig != nil { full := rs.FingersConfig.FullFingers idx = association.NewIndex() idx.BuildWithFingers(full.Fingers(), full.Aliases(), nil) } } + if idx == nil { + // cyberhub still answers, but without local fingerprint association. + d.Skip("cyberhub.index", deps.Name(engine.SetKey)+"/"+deps.Name(resources.SetKey)) + } cyberhub := NewCyberhubSearch(idx) reg.Register(commands.Command{Name: cyberhub.Name(), Usage: cyberhub.Usage(), Run: cyberhub.Run}, "search") }, diff --git a/tools/search/tavily.go b/tools/search/tavily.go index 339b3944..82c7b4b8 100644 --- a/tools/search/tavily.go +++ b/tools/search/tavily.go @@ -14,7 +14,7 @@ import ( "sync" "time" - "github.com/chainreactors/aiscan/pkg/agent/truncate" + "github.com/chainreactors/aiscan/core/truncate" ) const ( diff --git a/tools/search/websearch.go b/tools/search/websearch.go index 3dbef358..9f4ff5b8 100644 --- a/tools/search/websearch.go +++ b/tools/search/websearch.go @@ -4,7 +4,7 @@ import ( "fmt" "strings" - "github.com/chainreactors/aiscan/pkg/agent/provider" + "github.com/chainreactors/aiscan/agent/provider" ) func formatWebSearchResponse(resp *provider.WebSearchResponse, query string) string { diff --git a/tools/search/websearch_tool.go b/tools/search/websearch_tool.go index aad76dca..adec6c15 100644 --- a/tools/search/websearch_tool.go +++ b/tools/search/websearch_tool.go @@ -5,8 +5,8 @@ import ( "fmt" "strings" + "github.com/chainreactors/aiscan/agent/provider" "github.com/chainreactors/aiscan/core/tool" - "github.com/chainreactors/aiscan/pkg/agent/provider" ) type WebSearchTool struct { diff --git a/tools/spray/register.go b/tools/spray/register.go index 3489fdd2..fe478eee 100644 --- a/tools/spray/register.go +++ b/tools/spray/register.go @@ -1,22 +1,28 @@ package spray import ( - cfg "github.com/chainreactors/aiscan/core/config" + "github.com/chainreactors/aiscan/core/capability" + "github.com/chainreactors/aiscan/core/deps" "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/aiscan/tools/scan/engine" ) func init() { - cfg.ExtraScannerUsage["spray"] = func() string { return New(nil).Usage() } + capability.Register(capability.Descriptor{ + ID: "spray", Kind: capability.KindScanner, Group: "scanner", + CLIName: "spray", Summary: "spray", UsageLine: " spray Run spray directly", + Usage: func() string { return New(nil).Usage() }, Requires: []string{"scan.engine.Set.Spray"}, + }) commands.RegisterFactory(commands.Factory{ - Group: "scanner", - Build: func(deps *commands.Deps, reg *commands.CommandRegistry) { - es, _ := deps.EngineSet.(*engine.Set) - if es == nil || es.Spray == nil { + Capability: "spray", + Build: func(d *commands.Deps, reg *commands.CommandRegistry) { + es, ok := deps.Get(d.Bag, engine.SetKey) + if !ok || es == nil || es.Spray == nil { + d.Skip("spray", deps.Name(engine.SetKey)+".Spray") return } - impl := New(es.Spray).WithLogger(deps.GetLogger()).WithProxy(deps.ScannerProxy).WithDataBus(deps.DataBus) - reg.Register(commands.Command{Name: impl.Name(), Usage: impl.Usage(), QuickReference: impl.QuickReference(), Run: impl.Run, SetProxy: impl.SetProxy, GetProxy: func() string { return impl.Proxy }}, "scanner") + impl := New(es.Spray).WithLogger(d.GetLogger()).WithProxy(d.ScannerProxy).WithDataBus(d.DataBus) + reg.Register(commands.Command{Name: impl.Name(), Usage: impl.Usage(), QuickReference: impl.QuickReference(), Run: impl.Run}, "scanner") }, }) } diff --git a/tools/spray/spray.go b/tools/spray/spray.go index 69e0035c..142f016e 100644 --- a/tools/spray/spray.go +++ b/tools/spray/spray.go @@ -9,8 +9,8 @@ import ( "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" "github.com/chainreactors/aiscan/tools/toolargs" "github.com/chainreactors/sdk/spray" spraycore "github.com/chainreactors/spray/core" diff --git a/tools/spray/spray_test.go b/tools/spray/spray_test.go index 92091af2..91a2f89d 100644 --- a/tools/spray/spray_test.go +++ b/tools/spray/spray_test.go @@ -9,8 +9,8 @@ import ( "sync/atomic" "testing" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" sdkspray "github.com/chainreactors/sdk/spray" spraypkg "github.com/chainreactors/spray/pkg" "github.com/chainreactors/utils/parsers" diff --git a/tools/toolargs/base.go b/tools/toolargs/base.go index a83b0552..94a1e5a1 100644 --- a/tools/toolargs/base.go +++ b/tools/toolargs/base.go @@ -6,7 +6,7 @@ import ( "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" - "github.com/chainreactors/aiscan/pkg/telemetry" + "github.com/chainreactors/aiscan/core/telemetry" ) type Base struct { diff --git a/tools/zombie/register.go b/tools/zombie/register.go index 8b78f30b..90525aaa 100644 --- a/tools/zombie/register.go +++ b/tools/zombie/register.go @@ -1,22 +1,28 @@ package zombie import ( - cfg "github.com/chainreactors/aiscan/core/config" + "github.com/chainreactors/aiscan/core/capability" + "github.com/chainreactors/aiscan/core/deps" "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/aiscan/tools/scan/engine" ) func init() { - cfg.ExtraScannerUsage["zombie"] = func() string { return New(nil).Usage() } + capability.Register(capability.Descriptor{ + ID: "zombie", Kind: capability.KindScanner, Group: "scanner", + CLIName: "zombie", Summary: "zombie", UsageLine: " zombie Run zombie directly", + Usage: func() string { return New(nil).Usage() }, Requires: []string{"scan.engine.Set.Zombie"}, + }) commands.RegisterFactory(commands.Factory{ - Group: "scanner", - Build: func(deps *commands.Deps, reg *commands.CommandRegistry) { - es, _ := deps.EngineSet.(*engine.Set) - if es == nil || es.Zombie == nil { + Capability: "zombie", + Build: func(d *commands.Deps, reg *commands.CommandRegistry) { + es, ok := deps.Get(d.Bag, engine.SetKey) + if !ok || es == nil || es.Zombie == nil { + d.Skip("zombie", deps.Name(engine.SetKey)+".Zombie") return } - impl := New(es.Zombie).WithLogger(deps.GetLogger()).WithProxy(deps.ScannerProxy).WithDataBus(deps.DataBus) - reg.Register(commands.Command{Name: impl.Name(), Usage: impl.Usage(), Run: impl.Run, SetProxy: impl.SetProxy, GetProxy: func() string { return impl.Proxy }}, "scanner") + impl := New(es.Zombie).WithLogger(d.GetLogger()).WithProxy(d.ScannerProxy).WithDataBus(d.DataBus) + reg.Register(commands.Command{Name: impl.Name(), Usage: impl.Usage(), Run: impl.Run}, "scanner") }, }) } diff --git a/tools/zombie/zombie.go b/tools/zombie/zombie.go index 6d3e0966..00119323 100644 --- a/tools/zombie/zombie.go +++ b/tools/zombie/zombie.go @@ -8,8 +8,8 @@ import ( "github.com/chainreactors/aiscan/core/eventbus" "github.com/chainreactors/aiscan/core/output" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" "github.com/chainreactors/aiscan/tools/toolargs" sdkzombie "github.com/chainreactors/sdk/zombie" zombiecore "github.com/chainreactors/zombie/core" diff --git a/tools/zombie/zombie_test.go b/tools/zombie/zombie_test.go index 568d23c2..658ef0fc 100644 --- a/tools/zombie/zombie_test.go +++ b/tools/zombie/zombie_test.go @@ -9,8 +9,8 @@ import ( "strings" "testing" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/aiscan/pkg/telemetry" ) func TestExecuteDebugActivatesTelemetryLogger(t *testing.T) { diff --git a/web/frontend/cyber-ui b/web/frontend/cyber-ui index 1d1a29e1..2ab3d12a 160000 --- a/web/frontend/cyber-ui +++ b/web/frontend/cyber-ui @@ -1 +1 @@ -Subproject commit 1d1a29e17ecd3838880b193660df98c01cb1fc26 +Subproject commit 2ab3d12a1f312aacfd9319081fae6fd35a6733d8 diff --git a/web/frontend/e2e/aiscan-web.spec.ts b/web/frontend/e2e/aiscan-web.spec.ts index 58c71cf0..7b297178 100644 --- a/web/frontend/e2e/aiscan-web.spec.ts +++ b/web/frontend/e2e/aiscan-web.spec.ts @@ -1,4 +1,4 @@ -import { test, expect, type Page } from '@playwright/test'; +import { test, expect, type APIRequestContext, type Page } from '@playwright/test'; const API_TOKEN = process.env.ACCESS_KEY || 'test-token'; const LLM_PROVIDER = process.env.LLM_PROVIDER || 'openai'; @@ -19,6 +19,20 @@ async function openAuthenticatedApp(page: Page) { await expect(page.locator('button[aria-label="Open settings"]')).toBeVisible(); } +async function requireRegisteredAgents(request: APIRequestContext) { + let agents: any[] = []; + await expect.poll(async () => { + const response = await request.get('/api/agents', { headers: apiHeaders() }); + expect(response.ok()).toBeTruthy(); + agents = await response.json(); + return agents.length; + }, { + message: 'the E2E server must start and register its local mock-backed agent', + timeout: 15_000, + }).toBeGreaterThan(0); + return agents; +} + // --------------------------------------------------------------------------- // 1. Health & Status // --------------------------------------------------------------------------- @@ -302,15 +316,8 @@ test.describe('Agents API', () => { test.describe('Chat Session CRUD', () => { test('create, list, and delete a session', async ({ request }) => { // First, get available agents - const agentsRes = await request.get('/api/agents', { headers: apiHeaders() }); - const agents = await agentsRes.json(); - const agentID = agents.length > 0 ? agents[0].id : ''; - - // Skip if no agent is available - if (!agentID) { - test.skip(); - return; - } + const agents = await requireRegisteredAgents(request); + const agentID = agents[0].id; // Create const createRes = await request.post('/api/chat/sessions', { @@ -344,12 +351,7 @@ test.describe('Chat Session CRUD', () => { test.describe('Chat LLM round-trip', () => { test('send a message and receive an assistant response', async ({ request }) => { // Get agent - const agentsRes = await request.get('/api/agents', { headers: apiHeaders() }); - const agents = await agentsRes.json(); - if (agents.length === 0) { - test.skip(); - return; - } + const agents = await requireRegisteredAgents(request); const agentID = agents[0].id; // Create session @@ -375,7 +377,9 @@ test.describe('Chat LLM round-trip', () => { const msgRes = await request.get(`/api/chat/sessions/${sessionID}/messages`, { headers: apiHeaders(), }); - const messages = await msgRes.json(); + const page = await msgRes.json(); + expect(Array.isArray(page.items)).toBeTruthy(); + const messages = page.items; const assistantMsgs = messages.filter((m: any) => m.role === 'assistant'); if (assistantMsgs.length > 0) { assistantMsg = assistantMsgs[assistantMsgs.length - 1]; @@ -396,7 +400,44 @@ test.describe('Chat LLM round-trip', () => { }); // --------------------------------------------------------------------------- -// 10. SCO / Asset Pool API +// 10. SSE reconnect and durable event cursor +// --------------------------------------------------------------------------- + +test.describe('SSE reconnect', () => { + test('replays missing durable events after the browser reconnects', async ({ page, request, context }) => { + const agents = await requireRegisteredAgents(request); + + const createRes = await request.post('/api/chat/sessions', { + headers: { ...apiHeaders(), 'Content-Type': 'application/json' }, + data: { agent_id: agents[0].id }, + }); + expect(createRes.ok()).toBeTruthy(); + const session = await createRes.json(); + + await openAuthenticatedApp(page); + await page.goto(`/sessions/${session.id}`); + const prompt = 'Reply with exactly one word: PONG'; + const sendRes = await request.post(`/api/chat/sessions/${session.id}/messages`, { + headers: { ...apiHeaders(), 'Content-Type': 'application/json' }, + data: { content: prompt }, + }); + expect(sendRes.ok()).toBeTruthy(); + + await expect(page.locator('p').filter({ hasText: prompt })).toBeVisible({ timeout: 10_000 }); + await context.setOffline(true); + await page.waitForTimeout(3500); + await context.setOffline(false); + + const resumed = page.getByText('PONG', { exact: true }); + await expect(resumed).toBeVisible({ timeout: 15_000 }); + await expect(resumed).toHaveCount(1); + + await request.delete(`/api/chat/sessions/${session.id}`, { headers: apiHeaders() }); + }); +}); + +// --------------------------------------------------------------------------- +// 11. SCO / Asset Pool API // --------------------------------------------------------------------------- test.describe('Asset Pool API', () => { diff --git a/web/frontend/e2e/start-server.mjs b/web/frontend/e2e/start-server.mjs new file mode 100644 index 00000000..2d4fcf57 --- /dev/null +++ b/web/frontend/e2e/start-server.mjs @@ -0,0 +1,119 @@ +import { createServer } from 'node:http' +import { mkdtemp, rm, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join, resolve } from 'node:path' +import { fileURLToPath } from 'node:url' +import { spawn, spawnSync } from 'node:child_process' + +const host = '127.0.0.1' +const webPort = Number(process.env.AISCAN_E2E_PORT || 38080) +const root = resolve(fileURLToPath(new URL('../../..', import.meta.url))) +const workDir = await mkdtemp(join(tmpdir(), 'aiscan-web-e2e-')) +const binary = join(workDir, process.platform === 'win32' ? 'aiscan-e2e.exe' : 'aiscan-e2e') + +const mockLLM = createServer(async (req, res) => { + if (req.url === '/v1/models') { + res.writeHead(200, { 'Content-Type': 'application/json' }) + res.end(JSON.stringify({ data: [{ id: 'deepseek-chat', object: 'model' }] })) + return + } + if (req.url !== '/v1/chat/completions' || req.method !== 'POST') { + res.writeHead(404) + res.end('not found') + return + } + + const chunks = [] + for await (const chunk of req) chunks.push(chunk) + const payload = JSON.parse(Buffer.concat(chunks).toString('utf8') || '{}') + const delayedReply = JSON.stringify(payload.messages || []).includes('Reply with exactly one word: PONG') + if (payload.stream) { + res.writeHead(200, { + 'Content-Type': 'text/event-stream', + 'Cache-Control': 'no-cache', + Connection: 'keep-alive', + }) + if (delayedReply) { + res.write('data: {"choices":[{"delta":{"role":"assistant","content":"P"},"index":0}]}\n\n') + await new Promise((resolveDelay) => setTimeout(resolveDelay, 3000)) + res.write('data: {"choices":[{"delta":{"content":"ONG"},"index":0}]}\n\n') + } else { + res.write('data: {"choices":[{"delta":{"role":"assistant","content":"PONG"},"index":0}]}\n\n') + } + res.write('data: {"choices":[{"delta":{},"finish_reason":"stop","index":0}],"usage":{"prompt_tokens":10,"completion_tokens":1,"total_tokens":11}}\n\n') + res.end('data: [DONE]\n\n') + return + } + + if (delayedReply) await new Promise((resolveDelay) => setTimeout(resolveDelay, 3000)) + res.writeHead(200, { 'Content-Type': 'application/json' }) + res.end(JSON.stringify({ + id: 'chatcmpl-e2e', + choices: [{ message: { role: 'assistant', content: 'PONG' }, finish_reason: 'stop' }], + usage: { prompt_tokens: 10, completion_tokens: 1, total_tokens: 11 }, + })) +}) + +await new Promise((resolveListen, reject) => { + mockLLM.once('error', reject) + mockLLM.listen(0, host, resolveListen) +}) +const llmAddress = mockLLM.address() +if (!llmAddress || typeof llmAddress === 'string') throw new Error('mock LLM did not expose a TCP address') + +const configPath = join(workDir, 'aiscan.yaml') +await writeFile(configPath, `llm: + active_profile: e2e + providers: + - id: e2e + name: E2E DeepSeek + provider: deepseek + base_url: http://${host}:${llmAddress.port}/v1 + api_key: test-key + model: deepseek-chat +`, { mode: 0o600 }) + +const build = spawnSync('go', ['build', '-tags', 'full', '-o', binary, './cmd/aiscan'], { + cwd: root, + stdio: 'inherit', +}) +if (build.status !== 0) { + mockLLM.close() + await rm(workDir, { recursive: true, force: true }) + process.exit(build.status ?? 1) +} + +const child = spawn(binary, [ + '--config', configPath, + '--data-dir', join(workDir, 'data'), + 'web', + '--addr', `${host}:${webPort}`, + '--db', join(workDir, 'aiscan-web.db'), + '--token', 'test-token', +], { + cwd: root, + stdio: 'inherit', +}) + +let shuttingDown = false +async function shutdown(code) { + if (shuttingDown) return + shuttingDown = true + if (child.exitCode === null) child.kill() + await new Promise((resolveClose) => mockLLM.close(resolveClose)) + await rm(workDir, { recursive: true, force: true }) + process.exit(code) +} + +child.once('error', (error) => { + console.error(error) + void shutdown(1) +}) +child.once('exit', (code, signal) => { + if (!shuttingDown) { + console.error(`AIScan E2E server exited early (code=${code}, signal=${signal})`) + void shutdown(code ?? 1) + } +}) +process.once('SIGINT', () => void shutdown(130)) +process.once('SIGTERM', () => void shutdown(143)) diff --git a/web/frontend/playwright.config.ts b/web/frontend/playwright.config.ts index 1ca79740..293de43c 100644 --- a/web/frontend/playwright.config.ts +++ b/web/frontend/playwright.config.ts @@ -1,6 +1,7 @@ import { defineConfig } from '@playwright/test'; -const baseURL = process.env.BASE_URL || 'http://127.0.0.1:18080'; +const baseURL = process.env.BASE_URL || `http://127.0.0.1:${process.env.AISCAN_E2E_PORT || '38080'}`; +const manageServer = !process.env.BASE_URL; export default defineConfig({ testDir: './e2e', @@ -9,6 +10,14 @@ export default defineConfig({ fullyParallel: false, retries: 0, reporter: [['list'], ['html', { open: 'never' }]], + webServer: manageServer ? { + command: 'node ./e2e/start-server.mjs', + url: `${baseURL}/health`, + timeout: 180_000, + reuseExistingServer: false, + stdout: 'pipe', + stderr: 'pipe', + } : undefined, use: { baseURL, headless: true,