Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -187,7 +187,7 @@ aiscan agent --ioa-url http://127.0.0.1:8765 --space pentest-project \
export OPENAI_API_KEY="sk-..."

# CLI arguments
aiscan agent --provider deepseek --api-key sk-... --model deepseek-chat
aiscan agent --provider openai --base-url https://api.deepseek.com/v1 --api-key sk-... --model deepseek-chat
```

Config file `aiscan.yaml`:
Expand Down
2 changes: 1 addition & 1 deletion README_CN.md
Original file line number Diff line number Diff line change
Expand Up @@ -186,7 +186,7 @@ aiscan agent --ioa-url http://127.0.0.1:8765 --space pentest-project \
export OPENAI_API_KEY="sk-..."

# CLI 参数
aiscan agent --provider deepseek --api-key sk-... --model deepseek-chat
aiscan agent --provider openai --base-url https://api.deepseek.com/v1 --api-key sk-... --model deepseek-chat
```

配置文件 `aiscan.yaml`:
Expand Down
84 changes: 33 additions & 51 deletions agent/provider/provider.go
Original file line number Diff line number Diff line change
Expand Up @@ -42,33 +42,27 @@ type ProviderConfig struct {
ContextWindow int `yaml:"context_window,omitempty" config:"context_window"`
}

type providerPreset struct {
Protocol string
BaseURL string
APIKeyRequired bool
}
const (
ProviderOpenAI = "openai"
ProviderAnthropic = "anthropic"
)

var providerPresets = map[string]providerPreset{
"openai": {Protocol: "openai", BaseURL: "https://api.openai.com/v1", APIKeyRequired: true},
"anthropic": {Protocol: "anthropic", BaseURL: "https://api.anthropic.com/v1", APIKeyRequired: true},
"deepseek": {Protocol: "openai", BaseURL: "https://api.deepseek.com/v1", APIKeyRequired: true},
"openrouter": {Protocol: "openai", BaseURL: "https://openrouter.ai/api/v1", APIKeyRequired: true},
"groq": {Protocol: "openai", BaseURL: "https://api.groq.com/openai/v1", APIKeyRequired: true},
"moonshot": {Protocol: "openai", BaseURL: "https://api.moonshot.cn/v1", APIKeyRequired: true},
"ollama": {Protocol: "openai", BaseURL: "http://localhost:11434/v1"},
"zhipu": {Protocol: "openai", BaseURL: "https://open.bigmodel.cn/api/paas/v4", APIKeyRequired: true},
var providerBaseURLs = map[string]string{
ProviderOpenAI: "https://api.openai.com/v1",
ProviderAnthropic: "https://api.anthropic.com/v1",
}

var providerAliases = map[string]string{
"bigmodel": "zhipu",
"glm": "zhipu",
func NormalizeProvider(name string) string {
return strings.ToLower(strings.TrimSpace(name))
}

func NormalizeProvider(name string) string {
if strings.EqualFold(name, "anthropic") {
return "anthropic"
func IsSupportedProvider(name string) bool {
switch NormalizeProvider(name) {
case ProviderOpenAI, ProviderAnthropic:
return true
default:
return false
}
return "openai"
}

func Resolve(cfg *ProviderConfig) (*ProviderConfig, error) {
Expand All @@ -80,33 +74,19 @@ func Resolve(cfg *ProviderConfig) (*ProviderConfig, error) {
return nil, fmt.Errorf("context_window must be zero or positive")
}

providerName := strings.ToLower(strings.TrimSpace(resolved.Provider))
if alias, ok := providerAliases[providerName]; ok {
providerName = alias
}

providerName := NormalizeProvider(resolved.Provider)
if providerName == "" {
if resolved.BaseURL != "" {
providerName = InferFromBaseURL(resolved.BaseURL)
} else {
providerName = "openai"
}
providerName = InferFromBaseURL(resolved.BaseURL)
}

preset, knownProvider := providerPresets[providerName]
if knownProvider {
if strings.TrimSpace(resolved.BaseURL) == "" {
resolved.BaseURL = preset.BaseURL
}
resolved.Provider = preset.Protocol
} else {
if strings.TrimSpace(resolved.BaseURL) == "" {
return nil, fmt.Errorf("unknown provider %q: set base_url for a custom OpenAI-compatible endpoint", providerName)
}
resolved.Provider = NormalizeProvider(providerName)
if !IsSupportedProvider(providerName) {
return nil, fmt.Errorf("unsupported provider %q: use openai or anthropic", providerName)
}
if strings.TrimSpace(resolved.BaseURL) == "" {
resolved.BaseURL = providerBaseURLs[providerName]
}
resolved.Provider = providerName

if strings.TrimSpace(resolved.APIKey) == "" && (!knownProvider || preset.APIKeyRequired) {
if strings.TrimSpace(resolved.APIKey) == "" {
return nil, fmt.Errorf("no API key: set --api-key, llm.api_key, or AISCAN_API_KEY")
}

Expand All @@ -131,9 +111,7 @@ func NewProvider(cfg *ProviderConfig) (Provider, error) {
}

// inferImageSupport guesses whether a provider+model combination accepts
// image content parts based on the provider type and model name heuristics.
// Defaults to true for known provider types (anthropic/openai) and falls
// back to model-name heuristics for unknown providers.
// image content parts based on the protocol and model name heuristics.
func inferImageSupport(provider, model string) bool {
p := strings.ToLower(strings.TrimSpace(provider))
m := strings.ToLower(strings.TrimSpace(model))
Expand Down Expand Up @@ -163,16 +141,20 @@ func inferImageSupport(provider, model string) bool {
// silent failure.
func InferFromBaseURL(baseURL string) string {
if strings.Contains(strings.ToLower(baseURL), "anthropic.com") {
return "anthropic"
return ProviderAnthropic
}
return "openai"
return ProviderOpenAI
}

func NewProviderFromResolved(cfg *ProviderConfig) (Provider, error) {
if strings.ToLower(cfg.Provider) == "anthropic" {
switch NormalizeProvider(cfg.Provider) {
case ProviderAnthropic:
return NewAnthropicProvider(cfg)
case ProviderOpenAI:
return NewOpenAIProvider(cfg)
default:
return nil, fmt.Errorf("unsupported provider %q: use openai or anthropic", cfg.Provider)
}
return NewOpenAIProvider(cfg)
}

// Model capability registry extracted from pi's models.generated.ts.
Expand Down
36 changes: 8 additions & 28 deletions agent/provider/provider_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,14 +22,6 @@ func TestResolveProviderPresets(t *testing.T) {
}{
{name: "openai", provider: "openai", apiKey: "key", wantProtocol: "openai", wantBaseURL: "https://api.openai.com/v1"},
{name: "anthropic", provider: "anthropic", apiKey: "key", wantProtocol: "anthropic", wantBaseURL: "https://api.anthropic.com/v1"},
{name: "deepseek", provider: "deepseek", apiKey: "key", wantProtocol: "openai", wantBaseURL: "https://api.deepseek.com/v1"},
{name: "openrouter", provider: "openrouter", apiKey: "key", wantProtocol: "openai", wantBaseURL: "https://openrouter.ai/api/v1"},
{name: "groq", provider: "groq", apiKey: "key", wantProtocol: "openai", wantBaseURL: "https://api.groq.com/openai/v1"},
{name: "moonshot", provider: "moonshot", apiKey: "key", wantProtocol: "openai", wantBaseURL: "https://api.moonshot.cn/v1"},
{name: "ollama", provider: "ollama", wantProtocol: "openai", wantBaseURL: "http://localhost:11434/v1"},
{name: "zhipu", provider: "zhipu", apiKey: "key", wantProtocol: "openai", wantBaseURL: "https://open.bigmodel.cn/api/paas/v4"},
{name: "glm alias", provider: "glm", apiKey: "key", wantProtocol: "openai", wantBaseURL: "https://open.bigmodel.cn/api/paas/v4"},
{name: "bigmodel alias", provider: "bigmodel", apiKey: "key", wantProtocol: "openai", wantBaseURL: "https://open.bigmodel.cn/api/paas/v4"},
}

for _, tt := range tests {
Expand All @@ -45,30 +37,18 @@ func TestResolveProviderPresets(t *testing.T) {
}
}

func TestResolvePreservesExplicitDeepSeekBaseURL(t *testing.T) {
resolved, err := Resolve(&ProviderConfig{
Provider: "deepseek",
BaseURL: "https://gateway.example/v1",
APIKey: "key",
})
if err != nil {
t.Fatal(err)
}
if resolved.BaseURL != "https://gateway.example/v1" || resolved.Provider != "openai" {
t.Fatalf("Resolve() = %+v", resolved)
}
}

func TestResolveUnknownProviderRequiresBaseURL(t *testing.T) {
_, err := Resolve(&ProviderConfig{Provider: "custom", APIKey: "key"})
if err == nil || !strings.Contains(err.Error(), "base_url") {
t.Fatalf("Resolve() error = %v, want base_url guidance", err)
func TestResolveRejectsUnsupportedProvider(t *testing.T) {
for _, name := range []string{"deepseek", "openrouter", "ollama", "custom"} {
_, err := Resolve(&ProviderConfig{Provider: name, BaseURL: "https://gateway.example/v1", APIKey: "key"})
if err == nil || !strings.Contains(err.Error(), "use openai or anthropic") {
t.Fatalf("Resolve(%q) error = %v", name, err)
}
}
}

func TestResolveUsesBaseURL(t *testing.T) {
cfg, err := Resolve(&ProviderConfig{
Provider: "ollama",
Provider: "openai",
BaseURL: "http://localhost:11434/v1",
APIKey: "test-key",
})
Expand All @@ -93,7 +73,7 @@ func TestResolveRejectsNegativeModelLimits(t *testing.T) {

func TestResolvePreservesExplicitBaseURL(t *testing.T) {
cfg, err := Resolve(&ProviderConfig{
Provider: "ollama",
Provider: "openai",
BaseURL: "http://base-url.example/v1",
APIKey: "test-key",
})
Expand Down
13 changes: 9 additions & 4 deletions agent/retry.go
Original file line number Diff line number Diff line change
Expand Up @@ -96,11 +96,16 @@ func isRetryableByMessage(err error) bool {
// backward compatibility with external callers such as runner and webagent
// reconnect logic.
func RetryDelay(attempt int) time.Duration {
delay := time.Second << uint(attempt)
if delay > 10*time.Second {
delay = 10 * time.Second
if attempt < 0 {
attempt = 0
}
return delay
// Clamp before shifting. A large attempt previously overflowed the duration
// shift to zero, turning a persistent authentication failure into a tight
// reconnect loop that could saturate the control plane.
if attempt >= 4 {
return 10 * time.Second
}
return time.Second << uint(attempt)
}

// retryDelayFor computes the backoff for an LLM call retry. It honors a
Expand Down
7 changes: 7 additions & 0 deletions agent/retry_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -424,6 +424,13 @@ func TestRetryDelayBackoffSequence(t *testing.T) {
t.Errorf("attempt %d: RetryDelay = %s, want %s", i, got, w)
}
}

if got := RetryDelay(-1); got != time.Second {
t.Errorf("negative attempt: RetryDelay = %s, want 1s", got)
}
if got := RetryDelay(64); got != 10*time.Second {
t.Errorf("large attempt: RetryDelay = %s, want 10s", got)
}
}

func TestComputeRetryDelaySequence(t *testing.T) {
Expand Down
1 change: 1 addition & 0 deletions agent/types.go
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@ var (
ResolveProvider = provider.Resolve
InferProviderFromBaseURL = provider.InferFromBaseURL
NormalizeProvider = provider.NormalizeProvider
IsSupportedProvider = provider.IsSupportedProvider

ErrCallTimeout = provider.ErrCallTimeout
ErrStreamStalled = provider.ErrStreamStalled
Expand Down
6 changes: 3 additions & 3 deletions build.sh
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
# ./build.sh -o linux/amd64 # 快速编译单一平台
# ./build.sh -o "linux/amd64 darwin/arm64" # 编译指定平台
# ./build.sh --config prod.yaml # 使用指定配置文件
# ./build.sh --llm-model deepseek-chat # CLI 覆盖配置文件中的值
# ./build.sh --llm-provider openai --llm-model deepseek-chat # OpenAI-compatible
# ./build.sh --embed # 嵌入扫描资源(不加 emptytemplates/noembed tag)
# ./build.sh --ioa # 同时编译 ioa server 二进制

Expand Down Expand Up @@ -109,7 +109,7 @@ aiscan 构建脚本
--profile PROFILE 构建配置: agent (~28MB), mini (默认, ~77MB), full (~123MB)

LLM 覆盖(优先级高于 aiscan.yaml):
--llm-provider NAME
--llm-provider TYPE openai (OpenAI-compatible) or anthropic
--llm-base-url URL
--llm-api-key KEY
--llm-model NAME
Expand Down Expand Up @@ -137,7 +137,7 @@ Web Search:
./build.sh -o linux/amd64 # 快速编译单平台
./build.sh --config prod.yaml -o linux/amd64 # 使用生产配置编译
./build.sh --cyberhub-url http://10.0.0.1:9000 --cyberhub-key mykey
./build.sh --llm-provider deepseek --llm-model deepseek-chat
./build.sh --llm-provider openai --llm-base-url https://api.deepseek.com/v1 --llm-model deepseek-chat
./build.sh --embed # 嵌入资源的完整构建
./build.sh -g # 打印 ldflags(用于自定义构建命令)
./build.sh --profile agent -o linux/amd64 # agent 构建 (仅 agent REPL + Arsenal, 无内置扫描器)
Expand Down
8 changes: 4 additions & 4 deletions cmd/aiscan/cli_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -194,8 +194,8 @@ func TestParseCLIAgentAcceptsLLMFlags(t *testing.T) {
t.Fatalf("llm options = %#v", opt.LLMOptions)
}
pcfg := runner.ProviderConfig(&opt)
if pcfg.Provider != "" {
t.Fatalf("provider should be unresolved before agent.ResolveProvider, got %q", pcfg.Provider)
if pcfg.Provider != "openai" {
t.Fatalf("provider = %q, want openai protocol", pcfg.Provider)
}
resolved, err := agent.ResolveProvider(&pcfg)
if err != nil {
Expand Down Expand Up @@ -304,8 +304,8 @@ func TestParseCLIScanExtractsLLMFlags(t *testing.T) {
t.Fatalf("llm options = %#v", opt.LLMOptions)
}
pcfg := runner.ProviderConfig(&opt)
if pcfg.Provider != "" {
t.Fatalf("provider should be unresolved before agent.ResolveProvider, got %q", pcfg.Provider)
if pcfg.Provider != "openai" {
t.Fatalf("provider = %q, want openai protocol", pcfg.Provider)
}
resolved, err := agent.ResolveProvider(&pcfg)
if err != nil {
Expand Down
16 changes: 9 additions & 7 deletions cmd/aiscan/setup.go
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,7 @@ func initEngines(ctx context.Context, sc runner.ScannerConfig, logger telemetry.
HunterAPIKey: sc.HunterAPIKey,
IngressProxy: sc.ReconProxy,
Limit: sc.ReconLimit,
Credentials: sc.UncoverCredentials,
}
engineSet.SetupUncover(recon, logger)
return engineSet
Expand Down Expand Up @@ -97,13 +98,14 @@ func registerScannerCommands(cmdReg *commands.CommandRegistry, engineSet *engine

workDir, _ := os.Getwd()
deps := &commands.Deps{
WorkDir: workDir,
BashTimeout: toolCfg.BashTimeout,
SkillStore: skillStore,
ScannerProxy: scanCfg.Proxy,
Logger: logger,
TavilyKeys: toolCfg.TavilyKeys,
DataBus: dataBus,
WorkDir: workDir,
BashTimeout: toolCfg.BashTimeout,
SkillStore: skillStore,
ScannerProxy: scanCfg.Proxy,
Logger: logger,
TavilyKeys: toolCfg.TavilyKeys,
PlaywrightSession: toolCfg.PlaywrightSession,
DataBus: dataBus,
}
commands.Provide(deps, scan.OptsKey, scanOpts)
if engineSet != nil {
Expand Down
4 changes: 4 additions & 0 deletions cmd/aiscan/web_full.go
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,9 @@ func runWeb(ctx context.Context, option, explicitOption *cfg.Option, opts webCom
}
defer store.Close()

// The initial app must use the fully resolved option, including values loaded
// from the config file and environment. explicitOption is only the seed for
// later staged reloads, where the candidate config is resolved independently.
application, err := initWebApp(ctx, option, logger)
if err != nil {
return fmt.Errorf("init aiscan: %s", err)
Expand Down Expand Up @@ -269,6 +272,7 @@ func (s *webConfigStore) PrepareDistributeConfig(ctx context.Context, incoming w
}
current = parseDistributeConfig(data)
}
webproto.MigrateLLMConfig(&incoming.LLM, webproto.LLMProviderConfig{})

// Preserve existing secrets when incoming value is empty.
preserveLLMProfileSecrets(&incoming.LLM, current.LLM)
Expand Down
12 changes: 7 additions & 5 deletions core/config/config_gen.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,24 +9,26 @@ import (
const configFileHeader = `# aiscan 配置文件
#
# 运行时: aiscan 自动加载 ./aiscan.yaml 或 <二进制所在目录>/aiscan.yaml
# 优先级: CLI 参数 > 环境变量 > 配置文件 > 默认值
# 优先级: CLI > AIScan/集成环境变量 > 配置文件 > 协议环境变量 > 默认值
# 生成: aiscan --init
#
# 仅填写需要的字段,留空或删除的字段不会覆盖其他来源的值
#
# LLM 配置支持两种格式:
# 格式一 — 单 provider 简写(兼容旧配置):
# 格式一 — 单 provider 简写:
# llm:
# provider: deepseek
# provider: openai
# base_url: https://api.deepseek.com/v1
# api_key: sk-...
# model: deepseek-chat
#
# 格式二 — providers 配置列表(通过 active_profile 显式选择):
# 格式二 — LLM profile 列表(字段名 providers,通过 active_profile 选择):
# llm:
# active_profile: deepseek
# providers:
# - id: deepseek
# provider: deepseek
# provider: openai
# base_url: https://api.deepseek.com/v1
# api_key: sk-...
# model: deepseek-chat
# - id: openai
Expand Down
Loading