Merge branch 'main' of https://github.com/sipeed/picoclaw into feat-azure-openai-support
This commit is contained in:
commit
ec16f7cf8b
37 changed files with 3470 additions and 291 deletions
|
|
@ -5,6 +5,7 @@
|
|||
# ANTHROPIC_API_KEY=sk-ant-xxx
|
||||
# OPENAI_API_KEY=sk-xxx
|
||||
# GEMINI_API_KEY=xxx
|
||||
# MODELSCOPE_API_KEY=xxx
|
||||
# CLAUDE_CODE_OAUTH=xxx
|
||||
# ── Chat Channel ──────────────────────────
|
||||
# TELEGRAM_BOT_TOKEN=123456:ABC...
|
||||
|
|
|
|||
|
|
@ -27,6 +27,7 @@ builds:
|
|||
- windows
|
||||
- darwin
|
||||
- freebsd
|
||||
- netbsd
|
||||
goarch:
|
||||
- amd64
|
||||
- arm64
|
||||
|
|
@ -44,6 +45,12 @@ builds:
|
|||
ignore:
|
||||
- goos: windows
|
||||
goarch: arm
|
||||
- goos: netbsd
|
||||
goarch: s390x
|
||||
- goos: netbsd
|
||||
goarch: mips64
|
||||
- goos: netbsd
|
||||
goarch: arm
|
||||
|
||||
- id: picoclaw-launcher
|
||||
binary: picoclaw-launcher
|
||||
|
|
|
|||
2
Makefile
2
Makefile
|
|
@ -181,6 +181,8 @@ build-all: generate
|
|||
GOOS=linux GOARCH=arm GOARM=7 $(GO) build $(LDFLAGS) -o $(BUILD_DIR)/$(BINARY_NAME)-linux-armv7 ./$(CMD_DIR)
|
||||
GOOS=darwin GOARCH=arm64 $(GO) build $(LDFLAGS) -o $(BUILD_DIR)/$(BINARY_NAME)-darwin-arm64 ./$(CMD_DIR)
|
||||
GOOS=windows GOARCH=amd64 $(GO) build $(LDFLAGS) -o $(BUILD_DIR)/$(BINARY_NAME)-windows-amd64.exe ./$(CMD_DIR)
|
||||
GOOS=netbsd GOARCH=amd64 $(GO) build $(LDFLAGS) -o $(BUILD_DIR)/$(BINARY_NAME)-netbsd-amd64 ./$(CMD_DIR)
|
||||
GOOS=netbsd GOARCH=arm64 $(GO) build $(LDFLAGS) -o $(BUILD_DIR)/$(BINARY_NAME)-netbsd-arm64 ./$(CMD_DIR)
|
||||
@echo "All builds complete"
|
||||
|
||||
## install: Install picoclaw to system and copy builtin skills
|
||||
|
|
|
|||
|
|
@ -985,6 +985,7 @@ Cette conception permet également le **support multi-agent** avec une sélectio
|
|||
| **ShengsuanYun** | `shengsuanyun/` | `https://router.shengsuanyun.com/api/v1` | OpenAI | - |
|
||||
| **BytePlus** | `byteplus/` | `https://ark.ap-southeast.bytepluses.com/api/v3` | OpenAI | [Obtenir Clé](https://www.byteplus.com/) |
|
||||
| **LongCat** | `longcat/` | `https://api.longcat.chat/openai` | OpenAI | [Obtenir une clé](https://longcat.chat/platform) |
|
||||
| **ModelScope (魔搭)**| `modelscope/` | `https://api-inference.modelscope.cn/v1` | OpenAI | [Obtenir un Token](https://modelscope.cn/my/tokens) |
|
||||
| **Azure OpenAI** | `azure/` | `https://{resource}.openai.azure.com` | Azure | [Obtenir Clé](https://portal.azure.com) |
|
||||
| **Antigravity** | `antigravity/` | Google Cloud | Custom | OAuth uniquement |
|
||||
| **GitHub Copilot** | `github-copilot/` | `localhost:4321` | gRPC | - |
|
||||
|
|
@ -1224,6 +1225,7 @@ Cela se produit lorsqu'une autre instance du bot est en cours d'exécution. Assu
|
|||
| **Zhipu** | 200K tokens/mois | Convient aux utilisateurs chinois |
|
||||
| **Brave Search** | 2000 requêtes/mois | Fonctionnalité de recherche web |
|
||||
| **Groq** | Offre gratuite dispo | Inférence ultra-rapide (Llama, Mixtral) |
|
||||
| **ModelScope** | 2000 requêtes/jour | Inférence gratuite (Qwen, GLM, DeepSeek, etc.) |
|
||||
|
||||
---
|
||||
|
||||
|
|
|
|||
|
|
@ -926,6 +926,7 @@ HEARTBEAT_OK 応答 ユーザーが直接結果を受け取る
|
|||
| **ShengsuanYun** | `shengsuanyun/` | `https://router.shengsuanyun.com/api/v1` | OpenAI | - |
|
||||
| **BytePlus** | `byteplus/` | `https://ark.ap-southeast.bytepluses.com/api/v3` | OpenAI | [キーを取得](https://www.byteplus.com) |
|
||||
| **LongCat** | `longcat/` | `https://api.longcat.chat/openai` | OpenAI | [キーを取得](https://longcat.chat/platform) |
|
||||
| **ModelScope (魔搭)**| `modelscope/` | `https://api-inference.modelscope.cn/v1` | OpenAI | [トークンを取得](https://modelscope.cn/my/tokens) |
|
||||
| **Azure OpenAI** | `azure/` | `https://{resource}.openai.azure.com` | Azure | [キーを取得](https://portal.azure.com) |
|
||||
| **Antigravity** | `antigravity/` | Google Cloud | カスタム | OAuthのみ |
|
||||
| **GitHub Copilot** | `github-copilot/` | `localhost:4321` | gRPC | - |
|
||||
|
|
@ -1147,6 +1148,7 @@ Web 検索を有効にするには:
|
|||
| **Tavily** | 月 1000 クエリ | AI エージェント検索最適化 |
|
||||
| **Groq** | 無料枠あり | 高速推論(Llama, Mixtral) |
|
||||
| **Cerebras** | 無料枠あり | 高速推論(Llama, Qwen など) |
|
||||
| **ModelScope** | 1 日 2000 リクエスト | 無料推論(Qwen, GLM, DeepSeek など) |
|
||||
|
||||
---
|
||||
|
||||
|
|
|
|||
22
README.md
22
README.md
|
|
@ -1040,6 +1040,7 @@ This design also enables **multi-agent support** with flexible provider selectio
|
|||
| **BytePlus** | `byteplus/` | `https://ark.ap-southeast.bytepluses.com/api/v3` | OpenAI | [Get Key](https://www.byteplus.com) |
|
||||
| **Vivgrid** | `vivgrid/` | `https://api.vivgrid.com/v1` | OpenAI | [Get Key](https://vivgrid.com) |
|
||||
| **LongCat** | `longcat/` | `https://api.longcat.chat/openai` | OpenAI | [Get Key](https://longcat.chat/platform) |
|
||||
| **ModelScope (魔搭)**| `modelscope/` | `https://api-inference.modelscope.cn/v1` | OpenAI | [Get Token](https://modelscope.cn/my/tokens) |
|
||||
| **Azure OpenAI** | `azure/` | `https://{resource}.openai.azure.com` | Azure | [Get Key](https://portal.azure.com) |
|
||||
| **Antigravity** | `antigravity/` | Google Cloud | Custom | OAuth only |
|
||||
| **GitHub Copilot** | `github-copilot/` | `localhost:4321` | gRPC | - |
|
||||
|
|
@ -1132,6 +1133,26 @@ This design also enables **multi-agent support** with flexible provider selectio
|
|||
|
||||
> Run `picoclaw auth login --provider anthropic` to paste your API token.
|
||||
|
||||
**Anthropic Messages API (native format)**
|
||||
|
||||
For direct Anthropic API access or custom endpoints that only support Anthropic's native message format:
|
||||
|
||||
```json
|
||||
{
|
||||
"model_name": "claude-opus-4-6",
|
||||
"model": "anthropic-messages/claude-opus-4-6",
|
||||
"api_key": "sk-ant-your-key",
|
||||
"api_base": "https://api.anthropic.com"
|
||||
}
|
||||
```
|
||||
|
||||
> Use `anthropic-messages` protocol when:
|
||||
> - Using third-party proxies that only support Anthropic's native `/v1/messages` endpoint (not OpenAI-compatible `/v1/chat/completions`)
|
||||
> - Connecting to services like MiniMax, Synthetic that require Anthropic's native message format
|
||||
> - The existing `anthropic` protocol returns 404 errors (indicating the endpoint doesn't support OpenAI-compatible format)
|
||||
>
|
||||
> **Note:** The `anthropic` protocol uses OpenAI-compatible format (`/v1/chat/completions`), while `anthropic-messages` uses Anthropic's native format (`/v1/messages`). Choose based on your endpoint's supported format.
|
||||
|
||||
**Ollama (local)**
|
||||
|
||||
```json
|
||||
|
|
@ -1528,6 +1549,7 @@ This happens when another instance of the bot is running. Make sure only one `pi
|
|||
| **Groq** | Free tier available | Fast inference (Llama, Mixtral) |
|
||||
| **Cerebras** | Free tier available | Fast inference (Llama, Qwen, etc.) |
|
||||
| **LongCat** | Up to 5M tokens/day | Fast inference (free tier) |
|
||||
| **ModelScope** | 2000 requests/day | Free inference (Qwen, GLM, DeepSeek, etc.) |
|
||||
|
||||
---
|
||||
|
||||
|
|
|
|||
|
|
@ -981,6 +981,7 @@ Este design também possibilita o **suporte multi-agent** com seleção flexíve
|
|||
| **ShengsuanYun** | `shengsuanyun/` | `https://router.shengsuanyun.com/api/v1` | OpenAI | - |
|
||||
| **BytePlus** | `byteplus/` | `https://ark.ap-southeast.bytepluses.com/api/v3` | OpenAI | [Obter Chave](https://www.byteplus.com) |
|
||||
| **LongCat** | `longcat/` | `https://api.longcat.chat/openai` | OpenAI | [Obter Chave](https://longcat.chat/platform) |
|
||||
| **ModelScope (魔搭)**| `modelscope/` | `https://api-inference.modelscope.cn/v1` | OpenAI | [Obter Token](https://modelscope.cn/my/tokens) |
|
||||
| **Azure OpenAI** | `azure/` | `https://{resource}.openai.azure.com` | Azure | [Obter Chave](https://portal.azure.com) |
|
||||
| **Antigravity** | `antigravity/` | Google Cloud | Custom | Apenas OAuth |
|
||||
| **GitHub Copilot** | `github-copilot/` | `localhost:4321` | gRPC | - |
|
||||
|
|
@ -1221,6 +1222,7 @@ Isso acontece quando outra instância do bot está em execução. Certifique-se
|
|||
| **Brave Search** | 2000 consultas/mês | Funcionalidade de busca web |
|
||||
| **Groq** | Plano gratuito disponível | Inferência ultra-rápida (Llama, Mixtral) |
|
||||
| **Cerebras** | Plano gratuito disponível | Inferência ultra-rápida (Llama 3.3 70B) |
|
||||
| **ModelScope** | 2000 requisições/dia | Inferência gratuita (Qwen, GLM, DeepSeek, etc.) |
|
||||
|
||||
---
|
||||
|
||||
|
|
|
|||
|
|
@ -950,6 +950,7 @@ Thiết kế này cũng cho phép **hỗ trợ đa tác nhân** với lựa ch
|
|||
| **ShengsuanYun** | `shengsuanyun/` | `https://router.shengsuanyun.com/api/v1` | OpenAI | - |
|
||||
| **BytePlus** | `byteplus/` | `https://ark.ap-southeast.bytepluses.com/api/v3` | OpenAI | [Lấy Khóa](https://www.byteplus.com) |
|
||||
| **LongCat** | `longcat/` | `https://api.longcat.chat/openai` | OpenAI | [Lấy Key](https://longcat.chat/platform) |
|
||||
| **ModelScope (魔搭)**| `modelscope/` | `https://api-inference.modelscope.cn/v1` | OpenAI | [Lấy Token](https://modelscope.cn/my/tokens) |
|
||||
| **Azure OpenAI** | `azure/` | `https://{resource}.openai.azure.com` | Azure | [Lấy Khóa](https://portal.azure.com) |
|
||||
| **Antigravity** | `antigravity/` | Google Cloud | Tùy chỉnh | Chỉ OAuth |
|
||||
| **GitHub Copilot** | `github-copilot/` | `localhost:4321` | gRPC | - |
|
||||
|
|
@ -1189,6 +1190,7 @@ Một số nhà cung cấp (như Zhipu) có bộ lọc nội dung nghiêm ngặt
|
|||
| **Zhipu** | 200K tokens/tháng | Phù hợp cho người dùng Trung Quốc |
|
||||
| **Brave Search** | 2000 truy vấn/tháng | Chức năng tìm kiếm web |
|
||||
| **Groq** | Có gói miễn phí | Suy luận siêu nhanh (Llama, Mixtral) |
|
||||
| **ModelScope** | 2000 yêu cầu/ngày | Suy luận miễn phí (Qwen, GLM, DeepSeek, v.v.) |
|
||||
|
||||
---
|
||||
|
||||
|
|
|
|||
22
README.zh.md
22
README.zh.md
|
|
@ -522,6 +522,7 @@ Agent 读取 HEARTBEAT.md
|
|||
| **神算云** | `shengsuanyun/` | `https://router.shengsuanyun.com/api/v1` | OpenAI | - |
|
||||
| **BytePlus** | `byteplus/` | `https://ark.ap-southeast.bytepluses.com/api/v3` | OpenAI | [获取密钥](https://www.byteplus.com) |
|
||||
| **LongCat** | `longcat/` | `https://api.longcat.chat/openai` | OpenAI | [获取密钥](https://longcat.chat/platform) |
|
||||
| **ModelScope (魔搭)**| `modelscope/` | `https://api-inference.modelscope.cn/v1` | OpenAI | [获取 Token](https://modelscope.cn/my/tokens) |
|
||||
| **Azure OpenAI** | `azure/` | `https://{resource}.openai.azure.com` | Azure | [获取密钥](https://portal.azure.com) |
|
||||
| **Antigravity** | `antigravity/` | Google Cloud | 自定义 | 仅 OAuth |
|
||||
| **GitHub Copilot** | `github-copilot/` | `localhost:4321` | gRPC | - |
|
||||
|
|
@ -614,6 +615,26 @@ Agent 读取 HEARTBEAT.md
|
|||
|
||||
> 运行 `picoclaw auth login --provider anthropic` 来设置 OAuth 凭证。
|
||||
|
||||
**Anthropic Messages API(原生格式)**
|
||||
|
||||
用于直接访问 Anthropic API 或仅支持 Anthropic 原生消息格式的自定义端点:
|
||||
|
||||
```json
|
||||
{
|
||||
"model_name": "claude-opus-4-6",
|
||||
"model": "anthropic-messages/claude-opus-4-6",
|
||||
"api_key": "sk-ant-your-key",
|
||||
"api_base": "https://api.anthropic.com"
|
||||
}
|
||||
```
|
||||
|
||||
> 使用 `anthropic-messages` 协议的场景:
|
||||
> - 使用仅支持 Anthropic 原生 `/v1/messages` 端点的第三方代理(不支持 OpenAI 兼容的 `/v1/chat/completions`)
|
||||
> - 连接到 MiniMax、Synthetic 等需要 Anthropic 原生消息格式的服务
|
||||
> - 现有的 `anthropic` 协议返回 404 错误(说明端点不支持 OpenAI 兼容格式)
|
||||
>
|
||||
> **注意:** `anthropic` 协议使用 OpenAI 兼容格式(`/v1/chat/completions`),而 `anthropic-messages` 使用 Anthropic 原生格式(`/v1/messages`)。请根据端点支持的格式选择。
|
||||
|
||||
**Ollama (本地)**
|
||||
|
||||
```json
|
||||
|
|
@ -902,6 +923,7 @@ Discord: [https://discord.gg/V4sAZ9XWpN](https://discord.gg/V4sAZ9XWpN)
|
|||
| **Tavily** | 1000 次查询/月 | AI Agent 搜索优化 |
|
||||
| **Groq** | 提供免费层级 | 极速推理 (Llama, Mixtral) |
|
||||
| **LongCat** | 最多 5M tokens/天 | 推理速度快 (免费额度) |
|
||||
| **ModelScope (魔搭)** | 2000 次请求/天 | 免费推理 (Qwen, GLM, DeepSeek 等) |
|
||||
|
||||
---
|
||||
|
||||
|
|
|
|||
|
|
@ -3,10 +3,10 @@ package gateway
|
|||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
"os/signal"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/cmd/picoclaw/internal"
|
||||
|
|
@ -41,12 +41,31 @@ import (
|
|||
"github.com/sipeed/picoclaw/pkg/voice"
|
||||
)
|
||||
|
||||
// Timeout constants for service operations
|
||||
const (
|
||||
serviceRestartTimeout = 30 * time.Second
|
||||
serviceShutdownTimeout = 30 * time.Second
|
||||
providerReloadTimeout = 30 * time.Second
|
||||
gracefulShutdownTimeout = 15 * time.Second
|
||||
)
|
||||
|
||||
// gatewayServices holds references to all running services
|
||||
type gatewayServices struct {
|
||||
CronService *cron.CronService
|
||||
HeartbeatService *heartbeat.HeartbeatService
|
||||
MediaStore media.MediaStore
|
||||
ChannelManager *channels.Manager
|
||||
DeviceService *devices.Service
|
||||
HealthServer *health.Server
|
||||
}
|
||||
|
||||
func gatewayCmd(debug bool) error {
|
||||
if debug {
|
||||
logger.SetLevel(logger.DEBUG)
|
||||
fmt.Println("🔍 Debug mode enabled")
|
||||
}
|
||||
|
||||
configPath := internal.GetConfigPath()
|
||||
cfg, err := internal.LoadConfig()
|
||||
if err != nil {
|
||||
return fmt.Errorf("error loading config: %w", err)
|
||||
|
|
@ -83,9 +102,55 @@ func gatewayCmd(debug bool) error {
|
|||
"skills_available": skillsInfo["available"],
|
||||
})
|
||||
|
||||
// Setup and start all services
|
||||
services, err := setupAndStartServices(cfg, agentLoop, msgBus)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
fmt.Printf("✓ Gateway started on %s:%d\n", cfg.Gateway.Host, cfg.Gateway.Port)
|
||||
fmt.Println("Press Ctrl+C to stop")
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
go agentLoop.Run(ctx)
|
||||
|
||||
// Setup config file watcher for hot reload
|
||||
configReloadChan, stopWatch := setupConfigWatcherPolling(configPath, debug)
|
||||
defer stopWatch()
|
||||
|
||||
sigChan := make(chan os.Signal, 1)
|
||||
signal.Notify(sigChan, os.Interrupt)
|
||||
|
||||
// Main event loop - wait for signals or config changes
|
||||
for {
|
||||
select {
|
||||
case <-sigChan:
|
||||
logger.Info("Shutting down...")
|
||||
shutdownGateway(services, agentLoop, provider, true)
|
||||
return nil
|
||||
|
||||
case newCfg := <-configReloadChan:
|
||||
err := handleConfigReload(ctx, agentLoop, newCfg, &provider, services, msgBus)
|
||||
if err != nil {
|
||||
logger.Errorf("Config reload failed: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// setupAndStartServices initializes and starts all services
|
||||
func setupAndStartServices(
|
||||
cfg *config.Config,
|
||||
agentLoop *agent.AgentLoop,
|
||||
msgBus *bus.MessageBus,
|
||||
) (*gatewayServices, error) {
|
||||
services := &gatewayServices{}
|
||||
|
||||
// Setup cron tool and service
|
||||
execTimeout := time.Duration(cfg.Tools.Cron.ExecTimeoutMinutes) * time.Minute
|
||||
cronService := setupCronTool(
|
||||
services.CronService = setupCronTool(
|
||||
agentLoop,
|
||||
msgBus,
|
||||
cfg.WorkspacePath(),
|
||||
|
|
@ -93,20 +158,26 @@ func gatewayCmd(debug bool) error {
|
|||
execTimeout,
|
||||
cfg,
|
||||
)
|
||||
if err := services.CronService.Start(); err != nil {
|
||||
return nil, fmt.Errorf("error starting cron service: %w", err)
|
||||
}
|
||||
fmt.Println("✓ Cron service started")
|
||||
|
||||
heartbeatService := heartbeat.NewHeartbeatService(
|
||||
// Setup heartbeat service
|
||||
services.HeartbeatService = heartbeat.NewHeartbeatService(
|
||||
cfg.WorkspacePath(),
|
||||
cfg.Heartbeat.Interval,
|
||||
cfg.Heartbeat.Enabled,
|
||||
)
|
||||
heartbeatService.SetBus(msgBus)
|
||||
heartbeatService.SetHandler(func(prompt, channel, chatID string) *tools.ToolResult {
|
||||
services.HeartbeatService.SetBus(msgBus)
|
||||
services.HeartbeatService.SetHandler(func(prompt, channel, chatID string) *tools.ToolResult {
|
||||
// Use cli:direct as fallback if no valid channel
|
||||
if channel == "" || chatID == "" {
|
||||
channel, chatID = "cli", "direct"
|
||||
}
|
||||
// Use ProcessHeartbeat - no session history, each heartbeat is independent
|
||||
var response string
|
||||
var err error
|
||||
response, err = agentLoop.ProcessHeartbeat(context.Background(), prompt, channel, chatID)
|
||||
if err != nil {
|
||||
return tools.ErrorResult(fmt.Sprintf("Heartbeat error: %v", err))
|
||||
|
|
@ -118,24 +189,36 @@ func gatewayCmd(debug bool) error {
|
|||
// sent to user via processSystemMessage when the async task completes
|
||||
return tools.SilentResult(response)
|
||||
})
|
||||
if err := services.HeartbeatService.Start(); err != nil {
|
||||
return nil, fmt.Errorf("error starting heartbeat service: %w", err)
|
||||
}
|
||||
fmt.Println("✓ Heartbeat service started")
|
||||
|
||||
// Create media store for file lifecycle management with TTL cleanup
|
||||
mediaStore := media.NewFileMediaStoreWithCleanup(media.MediaCleanerConfig{
|
||||
services.MediaStore = media.NewFileMediaStoreWithCleanup(media.MediaCleanerConfig{
|
||||
Enabled: cfg.Tools.MediaCleanup.Enabled,
|
||||
MaxAge: time.Duration(cfg.Tools.MediaCleanup.MaxAge) * time.Minute,
|
||||
Interval: time.Duration(cfg.Tools.MediaCleanup.Interval) * time.Minute,
|
||||
})
|
||||
mediaStore.Start()
|
||||
// Start the media store if it's a FileMediaStore with cleanup
|
||||
if fms, ok := services.MediaStore.(*media.FileMediaStore); ok {
|
||||
fms.Start()
|
||||
}
|
||||
|
||||
channelManager, err := channels.NewManager(cfg, msgBus, mediaStore)
|
||||
// Create channel manager
|
||||
var err error
|
||||
services.ChannelManager, err = channels.NewManager(cfg, msgBus, services.MediaStore)
|
||||
if err != nil {
|
||||
mediaStore.Stop()
|
||||
return fmt.Errorf("error creating channel manager: %w", err)
|
||||
// Stop the media store if it's a FileMediaStore with cleanup
|
||||
if fms, ok := services.MediaStore.(*media.FileMediaStore); ok {
|
||||
fms.Stop()
|
||||
}
|
||||
return nil, fmt.Errorf("error creating channel manager: %w", err)
|
||||
}
|
||||
|
||||
// Inject channel manager and media store into agent loop
|
||||
agentLoop.SetChannelManager(channelManager)
|
||||
agentLoop.SetMediaStore(mediaStore)
|
||||
agentLoop.SetChannelManager(services.ChannelManager)
|
||||
agentLoop.SetMediaStore(services.MediaStore)
|
||||
|
||||
// Wire up voice transcription if a supported provider is configured.
|
||||
if transcriber := voice.DetectTranscriber(cfg); transcriber != nil {
|
||||
|
|
@ -143,83 +226,386 @@ func gatewayCmd(debug bool) error {
|
|||
logger.InfoCF("voice", "Transcription enabled (agent-level)", map[string]any{"provider": transcriber.Name()})
|
||||
}
|
||||
|
||||
enabledChannels := channelManager.GetEnabledChannels()
|
||||
enabledChannels := services.ChannelManager.GetEnabledChannels()
|
||||
if len(enabledChannels) > 0 {
|
||||
fmt.Printf("✓ Channels enabled: %s\n", enabledChannels)
|
||||
} else {
|
||||
fmt.Println("⚠ Warning: No channels enabled")
|
||||
}
|
||||
|
||||
fmt.Printf("✓ Gateway started on %s:%d\n", cfg.Gateway.Host, cfg.Gateway.Port)
|
||||
fmt.Println("Press Ctrl+C to stop")
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
if err := cronService.Start(); err != nil {
|
||||
fmt.Printf("Error starting cron service: %v\n", err)
|
||||
}
|
||||
fmt.Println("✓ Cron service started")
|
||||
|
||||
if err := heartbeatService.Start(); err != nil {
|
||||
fmt.Printf("Error starting heartbeat service: %v\n", err)
|
||||
}
|
||||
fmt.Println("✓ Heartbeat service started")
|
||||
|
||||
stateManager := state.NewManager(cfg.WorkspacePath())
|
||||
deviceService := devices.NewService(devices.Config{
|
||||
Enabled: cfg.Devices.Enabled,
|
||||
MonitorUSB: cfg.Devices.MonitorUSB,
|
||||
}, stateManager)
|
||||
deviceService.SetBus(msgBus)
|
||||
if err := deviceService.Start(ctx); err != nil {
|
||||
fmt.Printf("Error starting device service: %v\n", err)
|
||||
} else if cfg.Devices.Enabled {
|
||||
fmt.Println("✓ Device event service started")
|
||||
}
|
||||
|
||||
// Setup shared HTTP server with health endpoints and webhook handlers
|
||||
healthServer := health.NewServer(cfg.Gateway.Host, cfg.Gateway.Port)
|
||||
addr := fmt.Sprintf("%s:%d", cfg.Gateway.Host, cfg.Gateway.Port)
|
||||
channelManager.SetupHTTPServer(addr, healthServer)
|
||||
services.HealthServer = health.NewServer(cfg.Gateway.Host, cfg.Gateway.Port)
|
||||
services.ChannelManager.SetupHTTPServer(addr, services.HealthServer)
|
||||
|
||||
if err := channelManager.StartAll(ctx); err != nil {
|
||||
fmt.Printf("Error starting channels: %v\n", err)
|
||||
return err
|
||||
if err := services.ChannelManager.StartAll(context.Background()); err != nil {
|
||||
return nil, fmt.Errorf("error starting channels: %w", err)
|
||||
}
|
||||
|
||||
fmt.Printf("✓ Health endpoints available at http://%s:%d/health and /ready\n", cfg.Gateway.Host, cfg.Gateway.Port)
|
||||
|
||||
go agentLoop.Run(ctx)
|
||||
|
||||
sigChan := make(chan os.Signal, 1)
|
||||
signal.Notify(sigChan, os.Interrupt)
|
||||
<-sigChan
|
||||
|
||||
fmt.Println("\nShutting down...")
|
||||
if cp, ok := provider.(providers.StatefulProvider); ok {
|
||||
cp.Close()
|
||||
// Setup state manager and device service
|
||||
stateManager := state.NewManager(cfg.WorkspacePath())
|
||||
services.DeviceService = devices.NewService(devices.Config{
|
||||
Enabled: cfg.Devices.Enabled,
|
||||
MonitorUSB: cfg.Devices.MonitorUSB,
|
||||
}, stateManager)
|
||||
services.DeviceService.SetBus(msgBus)
|
||||
if err := services.DeviceService.Start(context.Background()); err != nil {
|
||||
logger.ErrorCF("device", "Error starting device service", map[string]any{"error": err.Error()})
|
||||
} else if cfg.Devices.Enabled {
|
||||
fmt.Println("✓ Device event service started")
|
||||
}
|
||||
cancel()
|
||||
msgBus.Close()
|
||||
|
||||
// Use a fresh context with timeout for graceful shutdown,
|
||||
// since the original ctx is already canceled.
|
||||
shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 15*time.Second)
|
||||
return services, nil
|
||||
}
|
||||
|
||||
// stopAndCleanupServices stops all services and cleans up resources
|
||||
func stopAndCleanupServices(
|
||||
services *gatewayServices,
|
||||
shutdownTimeout time.Duration,
|
||||
) {
|
||||
shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), shutdownTimeout)
|
||||
defer shutdownCancel()
|
||||
|
||||
channelManager.StopAll(shutdownCtx)
|
||||
deviceService.Stop()
|
||||
heartbeatService.Stop()
|
||||
cronService.Stop()
|
||||
mediaStore.Stop()
|
||||
if services.ChannelManager != nil {
|
||||
services.ChannelManager.StopAll(shutdownCtx)
|
||||
}
|
||||
if services.DeviceService != nil {
|
||||
services.DeviceService.Stop()
|
||||
}
|
||||
if services.HeartbeatService != nil {
|
||||
services.HeartbeatService.Stop()
|
||||
}
|
||||
if services.CronService != nil {
|
||||
services.CronService.Stop()
|
||||
}
|
||||
if services.MediaStore != nil {
|
||||
// Stop the media store if it's a FileMediaStore with cleanup
|
||||
if fms, ok := services.MediaStore.(*media.FileMediaStore); ok {
|
||||
fms.Stop()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// shutdownGateway performs a complete gateway shutdown
|
||||
func shutdownGateway(
|
||||
services *gatewayServices,
|
||||
agentLoop *agent.AgentLoop,
|
||||
provider providers.LLMProvider,
|
||||
fullShutdown bool,
|
||||
) {
|
||||
if cp, ok := provider.(providers.StatefulProvider); ok && fullShutdown {
|
||||
cp.Close()
|
||||
}
|
||||
|
||||
stopAndCleanupServices(services, gracefulShutdownTimeout)
|
||||
|
||||
agentLoop.Stop()
|
||||
agentLoop.Close()
|
||||
fmt.Println("✓ Gateway stopped")
|
||||
|
||||
logger.Info("✓ Gateway stopped")
|
||||
}
|
||||
|
||||
// handleConfigReload handles config file reload by stopping all services,
|
||||
// reloading the provider and config, and restarting services with the new config.
|
||||
func handleConfigReload(
|
||||
ctx context.Context,
|
||||
al *agent.AgentLoop,
|
||||
newCfg *config.Config,
|
||||
providerRef *providers.LLMProvider,
|
||||
services *gatewayServices,
|
||||
msgBus *bus.MessageBus,
|
||||
) error {
|
||||
logger.Info("🔄 Config file changed, reloading...")
|
||||
|
||||
newModel := newCfg.Agents.Defaults.ModelName
|
||||
if newModel == "" {
|
||||
newModel = newCfg.Agents.Defaults.Model
|
||||
}
|
||||
|
||||
logger.Infof(" New model is '%s', recreating provider...", newModel)
|
||||
|
||||
// Stop all services before reloading
|
||||
logger.Info(" Stopping all services...")
|
||||
stopAndCleanupServices(services, serviceShutdownTimeout)
|
||||
|
||||
// Create new provider from updated config first to ensure validity
|
||||
// This will use the correct API key and settings from newCfg.ModelList
|
||||
newProvider, newModelID, err := providers.CreateProvider(newCfg)
|
||||
if err != nil {
|
||||
logger.Errorf(" ⚠ Error creating new provider: %v", err)
|
||||
logger.Warn(" Attempting to restart services with old provider and config...")
|
||||
// Try to restart services with old configuration
|
||||
if restartErr := restartServices(al, services, msgBus); restartErr != nil {
|
||||
logger.Errorf(" ⚠ Failed to restart services: %v", restartErr)
|
||||
}
|
||||
return fmt.Errorf("error creating new provider: %w", err)
|
||||
}
|
||||
|
||||
if newModelID != "" {
|
||||
newCfg.Agents.Defaults.ModelName = newModelID
|
||||
}
|
||||
|
||||
// Use the atomic reload method on AgentLoop to safely swap provider and config.
|
||||
// This handles locking internally to prevent races with in-flight LLM calls
|
||||
// and concurrent reads of registry/config while the swap occurs.
|
||||
reloadCtx, reloadCancel := context.WithTimeout(context.Background(), providerReloadTimeout)
|
||||
defer reloadCancel()
|
||||
|
||||
if err := al.ReloadProviderAndConfig(reloadCtx, newProvider, newCfg); err != nil {
|
||||
logger.Errorf(" ⚠ Error reloading agent loop: %v", err)
|
||||
// Close the newly created provider since it wasn't adopted
|
||||
if cp, ok := newProvider.(providers.StatefulProvider); ok {
|
||||
cp.Close()
|
||||
}
|
||||
logger.Warn(" Attempting to restart services with old provider and config...")
|
||||
if restartErr := restartServices(al, services, msgBus); restartErr != nil {
|
||||
logger.Errorf(" ⚠ Failed to restart services: %v", restartErr)
|
||||
}
|
||||
return fmt.Errorf("error reloading agent loop: %w", err)
|
||||
}
|
||||
|
||||
// Update local provider reference only after successful atomic reload
|
||||
*providerRef = newProvider
|
||||
|
||||
// Restart all services with new config
|
||||
logger.Info(" Restarting all services with new configuration...")
|
||||
if err := restartServices(al, services, msgBus); err != nil {
|
||||
logger.Errorf(" ⚠ Error restarting services: %v", err)
|
||||
return fmt.Errorf("error restarting services: %w", err)
|
||||
}
|
||||
|
||||
logger.Info(" ✓ Provider, configuration, and services reloaded successfully (thread-safe)")
|
||||
return nil
|
||||
}
|
||||
|
||||
// restartServices restarts all services after a config reload
|
||||
func restartServices(
|
||||
al *agent.AgentLoop,
|
||||
services *gatewayServices,
|
||||
msgBus *bus.MessageBus,
|
||||
) error {
|
||||
// Create an independent context with timeout for service restart
|
||||
// This prevents cancellation from the main loop context during reload
|
||||
ctx, cancel := context.WithTimeout(context.Background(), serviceRestartTimeout)
|
||||
defer cancel()
|
||||
|
||||
// Get current config from agent loop (which has been updated if this is a reload)
|
||||
cfg := al.GetConfig()
|
||||
|
||||
// Re-create and start cron service with new config
|
||||
execTimeout := time.Duration(cfg.Tools.Cron.ExecTimeoutMinutes) * time.Minute
|
||||
services.CronService = setupCronTool(
|
||||
al,
|
||||
msgBus,
|
||||
cfg.WorkspacePath(),
|
||||
cfg.Agents.Defaults.RestrictToWorkspace,
|
||||
execTimeout,
|
||||
cfg,
|
||||
)
|
||||
if err := services.CronService.Start(); err != nil {
|
||||
return fmt.Errorf("error restarting cron service: %w", err)
|
||||
}
|
||||
fmt.Println(" ✓ Cron service restarted")
|
||||
|
||||
// Re-create and start heartbeat service with new config
|
||||
services.HeartbeatService = heartbeat.NewHeartbeatService(
|
||||
cfg.WorkspacePath(),
|
||||
cfg.Heartbeat.Interval,
|
||||
cfg.Heartbeat.Enabled,
|
||||
)
|
||||
services.HeartbeatService.SetBus(msgBus)
|
||||
services.HeartbeatService.SetHandler(func(prompt, channel, chatID string) *tools.ToolResult {
|
||||
if channel == "" || chatID == "" {
|
||||
channel, chatID = "cli", "direct"
|
||||
}
|
||||
var response string
|
||||
var err error
|
||||
response, err = al.ProcessHeartbeat(context.Background(), prompt, channel, chatID)
|
||||
if err != nil {
|
||||
return tools.ErrorResult(fmt.Sprintf("Heartbeat error: %v", err))
|
||||
}
|
||||
if response == "HEARTBEAT_OK" {
|
||||
return tools.SilentResult("Heartbeat OK")
|
||||
}
|
||||
return tools.SilentResult(response)
|
||||
})
|
||||
if err := services.HeartbeatService.Start(); err != nil {
|
||||
return fmt.Errorf("error restarting heartbeat service: %w", err)
|
||||
}
|
||||
fmt.Println(" ✓ Heartbeat service restarted")
|
||||
|
||||
// Stop the old media store before creating a new one
|
||||
if fms, ok := services.MediaStore.(*media.FileMediaStore); ok {
|
||||
fms.Stop()
|
||||
}
|
||||
|
||||
// Re-create media store with new config
|
||||
services.MediaStore = media.NewFileMediaStoreWithCleanup(media.MediaCleanerConfig{
|
||||
Enabled: cfg.Tools.MediaCleanup.Enabled,
|
||||
MaxAge: time.Duration(cfg.Tools.MediaCleanup.MaxAge) * time.Minute,
|
||||
Interval: time.Duration(cfg.Tools.MediaCleanup.Interval) * time.Minute,
|
||||
})
|
||||
// Start the media store if it's a FileMediaStore with cleanup
|
||||
if fms, ok := services.MediaStore.(*media.FileMediaStore); ok {
|
||||
fms.Start()
|
||||
}
|
||||
al.SetMediaStore(services.MediaStore)
|
||||
|
||||
// Re-create channel manager with new config
|
||||
var err error
|
||||
services.ChannelManager, err = channels.NewManager(cfg, msgBus, services.MediaStore)
|
||||
if err != nil {
|
||||
// Stop the media store if it's a FileMediaStore with cleanup
|
||||
if fms, ok := services.MediaStore.(*media.FileMediaStore); ok {
|
||||
fms.Stop()
|
||||
}
|
||||
return fmt.Errorf("error recreating channel manager: %w", err)
|
||||
}
|
||||
al.SetChannelManager(services.ChannelManager)
|
||||
|
||||
enabledChannels := services.ChannelManager.GetEnabledChannels()
|
||||
if len(enabledChannels) > 0 {
|
||||
fmt.Printf(" ✓ Channels enabled: %s\n", enabledChannels)
|
||||
} else {
|
||||
fmt.Println(" ⚠ Warning: No channels enabled")
|
||||
}
|
||||
|
||||
// Setup HTTP server with new config
|
||||
addr := fmt.Sprintf("%s:%d", cfg.Gateway.Host, cfg.Gateway.Port)
|
||||
services.HealthServer = health.NewServer(cfg.Gateway.Host, cfg.Gateway.Port)
|
||||
services.ChannelManager.SetupHTTPServer(addr, services.HealthServer)
|
||||
|
||||
if err := services.ChannelManager.StartAll(ctx); err != nil {
|
||||
return fmt.Errorf("error restarting channels: %w", err)
|
||||
}
|
||||
fmt.Printf(
|
||||
" ✓ Channels restarted, health endpoints at http://%s:%d/health and ready\n",
|
||||
cfg.Gateway.Host,
|
||||
cfg.Gateway.Port,
|
||||
)
|
||||
|
||||
// Re-create device service with new config
|
||||
stateManager := state.NewManager(cfg.WorkspacePath())
|
||||
services.DeviceService = devices.NewService(devices.Config{
|
||||
Enabled: cfg.Devices.Enabled,
|
||||
MonitorUSB: cfg.Devices.MonitorUSB,
|
||||
}, stateManager)
|
||||
services.DeviceService.SetBus(msgBus)
|
||||
if err := services.DeviceService.Start(ctx); err != nil {
|
||||
logger.WarnCF("device", "Failed to restart device service", map[string]any{"error": err.Error()})
|
||||
} else if cfg.Devices.Enabled {
|
||||
fmt.Println(" ✓ Device event service restarted")
|
||||
}
|
||||
|
||||
// Wire up voice transcription with new config
|
||||
transcriber := voice.DetectTranscriber(cfg)
|
||||
al.SetTranscriber(transcriber) // This will set it to nil if disabled
|
||||
if transcriber != nil {
|
||||
logger.InfoCF("voice", "Transcription re-enabled (agent-level)", map[string]any{"provider": transcriber.Name()})
|
||||
} else {
|
||||
logger.InfoCF("voice", "Transcription disabled", nil)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// setupConfigWatcherPolling sets up a simple polling-based config file watcher
|
||||
// Returns a channel for config updates and a stop function
|
||||
func setupConfigWatcherPolling(configPath string, debug bool) (chan *config.Config, func()) {
|
||||
configChan := make(chan *config.Config, 1)
|
||||
stop := make(chan struct{})
|
||||
var wg sync.WaitGroup
|
||||
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
|
||||
// Get initial file info
|
||||
lastModTime := getFileModTime(configPath)
|
||||
lastSize := getFileSize(configPath)
|
||||
|
||||
ticker := time.NewTicker(2 * time.Second) // Check every 2 seconds
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ticker.C:
|
||||
currentModTime := getFileModTime(configPath)
|
||||
currentSize := getFileSize(configPath)
|
||||
|
||||
// Check if file changed (modification time or size changed)
|
||||
if currentModTime.After(lastModTime) || currentSize != lastSize {
|
||||
if debug {
|
||||
logger.Debugf("🔍 Config file change detected")
|
||||
}
|
||||
|
||||
// Debounce - wait a bit to ensure file write is complete
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
|
||||
// Validate and load new config
|
||||
newCfg, err := config.LoadConfig(configPath)
|
||||
if err != nil {
|
||||
logger.Errorf("⚠ Error loading new config: %v", err)
|
||||
logger.Warn(" Using previous valid config")
|
||||
continue
|
||||
}
|
||||
|
||||
// Validate the new config
|
||||
if err := newCfg.ValidateModelList(); err != nil {
|
||||
logger.Errorf(" ⚠ New config validation failed: %v", err)
|
||||
logger.Warn(" Using previous valid config")
|
||||
continue
|
||||
}
|
||||
|
||||
logger.Info("✓ Config file validated and loaded")
|
||||
|
||||
// Update last known state
|
||||
lastModTime = currentModTime
|
||||
lastSize = currentSize
|
||||
|
||||
// Send new config to main loop (non-blocking)
|
||||
select {
|
||||
case configChan <- newCfg:
|
||||
default:
|
||||
// Channel full, skip this update
|
||||
logger.Warn("⚠ Previous config reload still in progress, skipping")
|
||||
}
|
||||
}
|
||||
|
||||
case <-stop:
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
stopFunc := func() {
|
||||
close(stop)
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
return configChan, stopFunc
|
||||
}
|
||||
|
||||
// getFileModTime returns the modification time of a file, or zero time if file doesn't exist
|
||||
func getFileModTime(path string) time.Time {
|
||||
info, err := os.Stat(path)
|
||||
if err != nil {
|
||||
return time.Time{}
|
||||
}
|
||||
return info.ModTime()
|
||||
}
|
||||
|
||||
// getFileSize returns the size of a file, or 0 if file doesn't exist
|
||||
func getFileSize(path string) int64 {
|
||||
info, err := os.Stat(path)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return info.Size()
|
||||
}
|
||||
|
||||
func setupCronTool(
|
||||
agentLoop *agent.AgentLoop,
|
||||
msgBus *bus.MessageBus,
|
||||
|
|
@ -239,7 +625,7 @@ func setupCronTool(
|
|||
var err error
|
||||
cronTool, err = tools.NewCronTool(cronService, agentLoop, msgBus, workspace, restrict, execTimeout, cfg)
|
||||
if err != nil {
|
||||
log.Fatalf("Critical error during CronTool initialization: %v", err)
|
||||
logger.Fatalf("Critical error during CronTool initialization: %v", err)
|
||||
}
|
||||
|
||||
agentLoop.RegisterTool(cronTool)
|
||||
|
|
|
|||
138
cmd/picoclaw/internal/model/command.go
Normal file
138
cmd/picoclaw/internal/model/command.go
Normal file
|
|
@ -0,0 +1,138 @@
|
|||
package model
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"github.com/sipeed/picoclaw/cmd/picoclaw/internal"
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
)
|
||||
|
||||
// LocalModel is a special model name that indicates that the model is local and with or without api_key.
|
||||
const LocalModel = "local-model"
|
||||
|
||||
func NewModelCommand() *cobra.Command {
|
||||
cmd := &cobra.Command{
|
||||
Use: "model [model_name]",
|
||||
Short: "Show or change the default model",
|
||||
Long: `Show or change the default model configuration.
|
||||
|
||||
If no argument is provided, shows the current default model.
|
||||
If a model name is provided, sets it as the default model.
|
||||
|
||||
Examples:
|
||||
picoclaw model # Show current default model
|
||||
picoclaw model gpt-5.2 # Set gpt-5.2 as default
|
||||
picoclaw model claude-sonnet-4.6 # Set claude-sonnet-4.6 as default
|
||||
picoclaw model local-model # Set local VLLM server as default
|
||||
|
||||
Note: 'local-model' is a special value for using a local VLLM server
|
||||
(running at localhost:8000 by default) which does not require an API key.`,
|
||||
Args: cobra.MaximumNArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
configPath := internal.GetConfigPath()
|
||||
|
||||
// Load current config
|
||||
cfg, err := config.LoadConfig(configPath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to load config: %w", err)
|
||||
}
|
||||
|
||||
if len(args) == 0 {
|
||||
// Show current default model
|
||||
showCurrentModel(cfg)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Set new default model
|
||||
modelName := args[0]
|
||||
return setDefaultModel(configPath, cfg, modelName)
|
||||
},
|
||||
}
|
||||
|
||||
return cmd
|
||||
}
|
||||
|
||||
func showCurrentModel(cfg *config.Config) {
|
||||
defaultModel := cfg.Agents.Defaults.ModelName
|
||||
if defaultModel == "" {
|
||||
defaultModel = cfg.Agents.Defaults.Model
|
||||
}
|
||||
|
||||
if defaultModel == "" {
|
||||
fmt.Println("No default model is currently set.")
|
||||
fmt.Println("\nAvailable models in your config:")
|
||||
listAvailableModels(cfg)
|
||||
} else {
|
||||
fmt.Printf("Current default model: %s\n", defaultModel)
|
||||
fmt.Println("\nAvailable models in your config:")
|
||||
listAvailableModels(cfg)
|
||||
}
|
||||
}
|
||||
|
||||
func listAvailableModels(cfg *config.Config) {
|
||||
if len(cfg.ModelList) == 0 {
|
||||
fmt.Println(" No models configured in model_list")
|
||||
return
|
||||
}
|
||||
|
||||
defaultModel := cfg.Agents.Defaults.ModelName
|
||||
if defaultModel == "" {
|
||||
defaultModel = cfg.Agents.Defaults.Model
|
||||
}
|
||||
|
||||
for _, model := range cfg.ModelList {
|
||||
marker := " "
|
||||
if model.ModelName == defaultModel {
|
||||
marker = "> "
|
||||
}
|
||||
if model.APIKey == "" {
|
||||
continue
|
||||
}
|
||||
fmt.Printf("%s- %s (%s)\n", marker, model.ModelName, model.Model)
|
||||
}
|
||||
}
|
||||
|
||||
func setDefaultModel(configPath string, cfg *config.Config, modelName string) error {
|
||||
// Validate that the model exists in model_list
|
||||
modelFound := false
|
||||
for _, model := range cfg.ModelList {
|
||||
if model.APIKey != "" && model.ModelName == modelName {
|
||||
modelFound = true
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if !modelFound && modelName != LocalModel {
|
||||
return fmt.Errorf("cannot found model '%s' in config", modelName)
|
||||
}
|
||||
|
||||
// Update the default model
|
||||
// Clear old model field and set new model_name
|
||||
oldModel := cfg.Agents.Defaults.ModelName
|
||||
if oldModel == "" {
|
||||
oldModel = cfg.Agents.Defaults.Model
|
||||
}
|
||||
|
||||
cfg.Agents.Defaults.ModelName = modelName
|
||||
cfg.Agents.Defaults.Model = "" // Clear deprecated field
|
||||
|
||||
// Save config back to file
|
||||
if err := config.SaveConfig(configPath, cfg); err != nil {
|
||||
return fmt.Errorf("failed to save config: %w", err)
|
||||
}
|
||||
|
||||
fmt.Printf("✓ Default model changed from '%s' to '%s'\n",
|
||||
formatModelName(oldModel), modelName)
|
||||
fmt.Println("\nThe new default model will be used for all agent interactions.")
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func formatModelName(name string) string {
|
||||
if name == "" {
|
||||
return "(none)"
|
||||
}
|
||||
return name
|
||||
}
|
||||
369
cmd/picoclaw/internal/model/command_test.go
Normal file
369
cmd/picoclaw/internal/model/command_test.go
Normal file
|
|
@ -0,0 +1,369 @@
|
|||
package model
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
)
|
||||
|
||||
var configPath = ""
|
||||
|
||||
func initTest(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
configPath = filepath.Join(tmpDir, "config.json")
|
||||
_ = os.Setenv("PICOCLAW_CONFIG", configPath)
|
||||
}
|
||||
|
||||
// captureStdout captures stdout during the execution of fn and returns the captured output
|
||||
func captureStdout(fn func()) string {
|
||||
oldStdout := os.Stdout
|
||||
r, w, _ := os.Pipe()
|
||||
os.Stdout = w
|
||||
|
||||
fn()
|
||||
|
||||
w.Close()
|
||||
os.Stdout = oldStdout
|
||||
|
||||
var buf bytes.Buffer
|
||||
io.Copy(&buf, r)
|
||||
return buf.String()
|
||||
}
|
||||
|
||||
func TestNewModelCommand(t *testing.T) {
|
||||
cmd := NewModelCommand()
|
||||
|
||||
require.NotNil(t, cmd)
|
||||
|
||||
assert.Equal(t, "model [model_name]", cmd.Use)
|
||||
assert.Equal(t, "Show or change the default model", cmd.Short)
|
||||
|
||||
assert.Len(t, cmd.Aliases, 0)
|
||||
|
||||
assert.False(t, cmd.HasFlags())
|
||||
|
||||
assert.Nil(t, cmd.Run)
|
||||
assert.NotNil(t, cmd.RunE)
|
||||
|
||||
assert.Nil(t, cmd.PersistentPreRunE)
|
||||
assert.Nil(t, cmd.PersistentPreRun)
|
||||
assert.Nil(t, cmd.PersistentPostRun)
|
||||
}
|
||||
|
||||
func TestShowCurrentModel_WithDefaultModel(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
ModelName: "gpt-4",
|
||||
},
|
||||
},
|
||||
ModelList: []config.ModelConfig{
|
||||
{ModelName: "gpt-4", Model: "openai/gpt-4", APIKey: "test"},
|
||||
{ModelName: "claude-3", Model: "anthropic/claude-3", APIKey: "test"},
|
||||
},
|
||||
}
|
||||
|
||||
output := captureStdout(func() {
|
||||
showCurrentModel(cfg)
|
||||
})
|
||||
|
||||
assert.Contains(t, output, "Current default model: gpt-4")
|
||||
assert.Contains(t, output, "Available models in your config:")
|
||||
assert.Contains(t, output, "gpt-4")
|
||||
assert.Contains(t, output, "claude-3")
|
||||
}
|
||||
|
||||
func TestShowCurrentModel_NoDefaultModel(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
ModelName: "",
|
||||
Model: "",
|
||||
},
|
||||
},
|
||||
ModelList: []config.ModelConfig{
|
||||
{ModelName: "gpt-4", Model: "openai/gpt-4", APIKey: "test"},
|
||||
},
|
||||
}
|
||||
|
||||
output := captureStdout(func() {
|
||||
showCurrentModel(cfg)
|
||||
})
|
||||
|
||||
assert.Contains(t, output, "No default model is currently set.")
|
||||
assert.Contains(t, output, "Available models in your config:")
|
||||
}
|
||||
|
||||
func TestShowCurrentModel_BackwardCompatibility(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
Model: "legacy-model",
|
||||
},
|
||||
},
|
||||
ModelList: []config.ModelConfig{},
|
||||
}
|
||||
|
||||
output := captureStdout(func() {
|
||||
showCurrentModel(cfg)
|
||||
})
|
||||
|
||||
assert.Contains(t, output, "Current default model: legacy-model")
|
||||
}
|
||||
|
||||
func TestListAvailableModels_Empty(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
ModelList: []config.ModelConfig{},
|
||||
}
|
||||
|
||||
output := captureStdout(func() {
|
||||
listAvailableModels(cfg)
|
||||
})
|
||||
|
||||
assert.Contains(t, output, "No models configured in model_list")
|
||||
}
|
||||
|
||||
func TestListAvailableModels_WithModels(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
ModelName: "gpt-4",
|
||||
},
|
||||
},
|
||||
ModelList: []config.ModelConfig{
|
||||
{ModelName: "gpt-4", Model: "openai/gpt-4", APIKey: "test"},
|
||||
{ModelName: "claude-3", Model: "anthropic/claude-3", APIKey: "test"},
|
||||
{ModelName: "no-key-model", Model: "openai/test", APIKey: ""},
|
||||
},
|
||||
}
|
||||
|
||||
output := captureStdout(func() {
|
||||
listAvailableModels(cfg)
|
||||
})
|
||||
|
||||
assert.NotEmpty(t, output)
|
||||
assert.Contains(t, output, "> - gpt-4 (openai/gpt-4)")
|
||||
assert.Contains(t, output, "claude-3 (anthropic/claude-3)")
|
||||
assert.NotContains(t, output, "no-key-model")
|
||||
}
|
||||
|
||||
func TestSetDefaultModel_ValidModel(t *testing.T) {
|
||||
initTest(t)
|
||||
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
ModelName: "old-model",
|
||||
},
|
||||
},
|
||||
ModelList: []config.ModelConfig{
|
||||
{ModelName: "new-model", Model: "openai/new-model", APIKey: "test"},
|
||||
{ModelName: "old-model", Model: "openai/old-model", APIKey: "test"},
|
||||
},
|
||||
}
|
||||
|
||||
output := captureStdout(func() {
|
||||
err := setDefaultModel(configPath, cfg, "new-model")
|
||||
assert.NoError(t, err)
|
||||
})
|
||||
|
||||
assert.Contains(t, output, "Default model changed from 'old-model' to 'new-model'")
|
||||
|
||||
// Verify config was updated
|
||||
updatedCfg, err := config.LoadConfig(configPath)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "new-model", updatedCfg.Agents.Defaults.ModelName)
|
||||
assert.Empty(t, updatedCfg.Agents.Defaults.Model)
|
||||
}
|
||||
|
||||
func TestSetDefaultModel_LegacyModelField(t *testing.T) {
|
||||
initTest(t)
|
||||
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
Model: "legacy-old",
|
||||
},
|
||||
},
|
||||
ModelList: []config.ModelConfig{
|
||||
{ModelName: "new-model", Model: "openai/new-model", APIKey: "test"},
|
||||
},
|
||||
}
|
||||
|
||||
output := captureStdout(func() {
|
||||
err := setDefaultModel(configPath, cfg, "new-model")
|
||||
assert.NoError(t, err)
|
||||
})
|
||||
|
||||
assert.Contains(t, output, "Default model changed from 'legacy-old' to 'new-model'")
|
||||
}
|
||||
|
||||
func TestSetDefaultModel_InvalidModel(t *testing.T) {
|
||||
initTest(t)
|
||||
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
ModelName: "existing-model",
|
||||
},
|
||||
},
|
||||
ModelList: []config.ModelConfig{
|
||||
{ModelName: "existing-model", Model: "openai/existing", APIKey: "test"},
|
||||
},
|
||||
}
|
||||
|
||||
assert.Error(t, setDefaultModel(configPath, cfg, "nonexistent-model"))
|
||||
}
|
||||
|
||||
func TestSetDefaultModel_ModelWithoutAPIKey(t *testing.T) {
|
||||
initTest(t)
|
||||
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
ModelName: "existing-model",
|
||||
},
|
||||
},
|
||||
ModelList: []config.ModelConfig{
|
||||
{ModelName: "existing-model", Model: "openai/existing", APIKey: "test"},
|
||||
{ModelName: "no-key-model", Model: "openai/nokey", APIKey: ""},
|
||||
},
|
||||
}
|
||||
|
||||
assert.Error(t, setDefaultModel(configPath, cfg, "no-key-model"))
|
||||
}
|
||||
|
||||
func TestSetDefaultModel_SaveConfigError(t *testing.T) {
|
||||
// Use an invalid path to trigger save error
|
||||
invalidPath := "/nonexistent/directory/config.json"
|
||||
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
ModelName: "old-model",
|
||||
},
|
||||
},
|
||||
ModelList: []config.ModelConfig{
|
||||
{ModelName: "new-model", Model: "openai/new-model", APIKey: "test"},
|
||||
},
|
||||
}
|
||||
|
||||
err := setDefaultModel(invalidPath, cfg, "new-model")
|
||||
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "failed to save config")
|
||||
}
|
||||
|
||||
func TestFormatModelName(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
expected string
|
||||
}{
|
||||
{"empty string", "", "(none)"},
|
||||
{"simple model", "gpt-4", "gpt-4"},
|
||||
{"model with version", "claude-sonnet-4.6", "claude-sonnet-4.6"},
|
||||
{"model with spaces", "my model", "my model"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := formatModelName(tt.input)
|
||||
assert.Equal(t, tt.expected, result)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelCommandExecution_Show(t *testing.T) {
|
||||
initTest(t)
|
||||
|
||||
// Create a test config
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
ModelName: "test-model",
|
||||
},
|
||||
},
|
||||
ModelList: []config.ModelConfig{
|
||||
{ModelName: "test-model", Model: "openai/test", APIKey: "test"},
|
||||
},
|
||||
}
|
||||
|
||||
err := config.SaveConfig(configPath, cfg)
|
||||
require.NoError(t, err)
|
||||
|
||||
cmd := NewModelCommand()
|
||||
|
||||
output := captureStdout(func() {
|
||||
err = cmd.RunE(cmd, []string{})
|
||||
assert.NoError(t, err)
|
||||
})
|
||||
|
||||
assert.Contains(t, output, "Current default model: test-model")
|
||||
}
|
||||
|
||||
func TestModelCommandExecution_Set(t *testing.T) {
|
||||
initTest(t)
|
||||
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
ModelName: "old-model",
|
||||
},
|
||||
},
|
||||
ModelList: []config.ModelConfig{
|
||||
{ModelName: "old-model", Model: "openai/old", APIKey: "test"},
|
||||
{ModelName: "new-model", Model: "openai/new", APIKey: "test"},
|
||||
},
|
||||
}
|
||||
|
||||
err := config.SaveConfig(configPath, cfg)
|
||||
require.NoError(t, err)
|
||||
|
||||
cmd := NewModelCommand()
|
||||
|
||||
output := captureStdout(func() {
|
||||
err = cmd.RunE(cmd, []string{"new-model"})
|
||||
assert.NoError(t, err)
|
||||
})
|
||||
|
||||
assert.Contains(t, output, "Default model changed from 'old-model' to 'new-model'")
|
||||
}
|
||||
|
||||
func TestModelCommandExecution_TooManyArgs(t *testing.T) {
|
||||
cmd := NewModelCommand()
|
||||
|
||||
err := cmd.RunE(cmd, []string{"model1", "model2"})
|
||||
|
||||
assert.Error(t, err)
|
||||
}
|
||||
|
||||
func TestListAvailableModels_MarkerLogic(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
ModelName: "middle-model",
|
||||
},
|
||||
},
|
||||
ModelList: []config.ModelConfig{
|
||||
{ModelName: "first-model", Model: "openai/first", APIKey: "test"},
|
||||
{ModelName: "middle-model", Model: "openai/middle", APIKey: "test"},
|
||||
{ModelName: "last-model", Model: "openai/last", APIKey: "test"},
|
||||
},
|
||||
}
|
||||
|
||||
output := captureStdout(func() {
|
||||
listAvailableModels(cfg)
|
||||
})
|
||||
|
||||
assert.Contains(t, output, " - first-model (openai/first)")
|
||||
assert.Contains(t, output, "> - middle-model (openai/middle)")
|
||||
assert.Contains(t, output, " - last-model (openai/last)")
|
||||
}
|
||||
|
|
@ -29,7 +29,15 @@ func NewSkillsCommand() *cobra.Command {
|
|||
}
|
||||
|
||||
d.workspace = cfg.WorkspacePath()
|
||||
d.installer = skills.NewSkillInstaller(d.workspace)
|
||||
installer, err := skills.NewSkillInstaller(
|
||||
d.workspace,
|
||||
cfg.Tools.Skills.Github.Token,
|
||||
cfg.Tools.Skills.Github.Proxy,
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("error creating skills installer: %w", err)
|
||||
}
|
||||
d.installer = installer
|
||||
|
||||
// get global config directory and builtin skills directory
|
||||
globalDir := filepath.Dir(internal.GetConfigPath())
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ import (
|
|||
"github.com/sipeed/picoclaw/cmd/picoclaw/internal/cron"
|
||||
"github.com/sipeed/picoclaw/cmd/picoclaw/internal/gateway"
|
||||
"github.com/sipeed/picoclaw/cmd/picoclaw/internal/migrate"
|
||||
"github.com/sipeed/picoclaw/cmd/picoclaw/internal/model"
|
||||
"github.com/sipeed/picoclaw/cmd/picoclaw/internal/onboard"
|
||||
"github.com/sipeed/picoclaw/cmd/picoclaw/internal/skills"
|
||||
"github.com/sipeed/picoclaw/cmd/picoclaw/internal/status"
|
||||
|
|
@ -43,6 +44,7 @@ func NewPicoclawCommand() *cobra.Command {
|
|||
cron.NewCronCommand(),
|
||||
migrate.NewMigrateCommand(),
|
||||
skills.NewSkillsCommand(),
|
||||
model.NewModelCommand(),
|
||||
version.NewVersionCommand(),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -39,6 +39,7 @@ func TestNewPicoclawCommand(t *testing.T) {
|
|||
"cron",
|
||||
"gateway",
|
||||
"migrate",
|
||||
"model",
|
||||
"onboard",
|
||||
"skills",
|
||||
"status",
|
||||
|
|
|
|||
|
|
@ -25,6 +25,13 @@
|
|||
"api_base": "https://api.anthropic.com/v1",
|
||||
"thinking_level": "high"
|
||||
},
|
||||
{
|
||||
"_comment": "Anthropic Messages API - use native format for direct Anthropic API access",
|
||||
"model_name": "claude-opus-4-6",
|
||||
"model": "anthropic-messages/claude-opus-4-6",
|
||||
"api_key": "sk-ant-your-key",
|
||||
"api_base": "https://api.anthropic.com"
|
||||
},
|
||||
{
|
||||
"model_name": "gemini",
|
||||
"model": "antigravity/gemini-2.0-flash",
|
||||
|
|
@ -40,6 +47,12 @@
|
|||
"model": "longcat/LongCat-Flash-Thinking",
|
||||
"api_key": "your-longcat-api-key"
|
||||
},
|
||||
{
|
||||
"model_name": "modelscope-qwen",
|
||||
"model": "modelscope/Qwen/Qwen3-235B-A22B-Instruct-2507",
|
||||
"api_key": "your-modelscope-access-token",
|
||||
"api_base": "https://api-inference.modelscope.cn/v1"
|
||||
},
|
||||
{
|
||||
"model_name": "azure-gpt5",
|
||||
"model": "azure/my-gpt5-deployment",
|
||||
|
|
@ -289,6 +302,10 @@
|
|||
"longcat": {
|
||||
"api_key": "",
|
||||
"api_base": "https://api.longcat.chat/openai"
|
||||
},
|
||||
"modelscope": {
|
||||
"api_key": "",
|
||||
"api_base": "https://api-inference.modelscope.cn/v1"
|
||||
}
|
||||
},
|
||||
"tools": {
|
||||
|
|
@ -433,6 +450,10 @@
|
|||
"max_response_size": 0
|
||||
}
|
||||
},
|
||||
"github": {
|
||||
"proxy": "http://127.0.0.1:7891",
|
||||
"token": ""
|
||||
},
|
||||
"max_concurrent_searches": 2,
|
||||
"search_cache": {
|
||||
"max_size": 50,
|
||||
|
|
|
|||
2
go.sum
2
go.sum
|
|
@ -271,8 +271,6 @@ golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug
|
|||
golang.org/x/net v0.6.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs=
|
||||
golang.org/x/net v0.10.0/go.mod h1:0qNGK6F8kojg2nk9dLZ2mShWaEBan6FAoqfSigmmuDg=
|
||||
golang.org/x/net v0.19.0/go.mod h1:CfAk/cbD4CthTvqiEl8NpboMuiuOYsAr/7NOjZJtv1U=
|
||||
golang.org/x/net v0.50.0 h1:ucWh9eiCGyDR3vtzso0WMQinm2Dnt8cFMuQa9K33J60=
|
||||
golang.org/x/net v0.50.0/go.mod h1:UgoSli3F/pBgdJBHCTc+tp3gmrU4XswgGRgtnwWTfyM=
|
||||
golang.org/x/net v0.51.0 h1:94R/GTO7mt3/4wIKpcR5gkGmRLOuE/2hNGeWq/GBIFo=
|
||||
golang.org/x/net v0.51.0/go.mod h1:aamm+2QF5ogm02fjy5Bb7CQ0WMt1/WVM7FtyaTLlA9Y=
|
||||
golang.org/x/oauth2 v0.23.0/go.mod h1:XYTD2NtWslqkgxebSiOHnXEap4TF09sJSc7H1sXbhtI=
|
||||
|
|
|
|||
|
|
@ -48,6 +48,9 @@ type AgentLoop struct {
|
|||
transcriber voice.Transcriber
|
||||
cmdRegistry *commands.Registry
|
||||
mcp mcpRuntime
|
||||
mu sync.RWMutex
|
||||
// Track active requests for safe provider cleanup
|
||||
activeRequests sync.WaitGroup
|
||||
}
|
||||
|
||||
// processOptions configures how a message is processed
|
||||
|
|
@ -239,6 +242,7 @@ func registerSharedTools(
|
|||
|
||||
func (al *AgentLoop) Run(ctx context.Context) error {
|
||||
al.running.Store(true)
|
||||
|
||||
if err := al.ensureMCPInitialized(ctx); err != nil {
|
||||
return err
|
||||
}
|
||||
|
|
@ -278,7 +282,7 @@ func (al *AgentLoop) Run(ctx context.Context) error {
|
|||
// If so, skip publishing to avoid duplicate messages to the user.
|
||||
// Use default agent's tools to check (message tool is shared).
|
||||
alreadySent := false
|
||||
defaultAgent := al.registry.GetDefaultAgent()
|
||||
defaultAgent := al.GetRegistry().GetDefaultAgent()
|
||||
if defaultAgent != nil {
|
||||
if tool, ok := defaultAgent.Tools.Get("message"); ok {
|
||||
if mt, ok := tool.(*tools.MessageTool); ok {
|
||||
|
|
@ -331,12 +335,13 @@ func (al *AgentLoop) Close() {
|
|||
}
|
||||
}
|
||||
|
||||
al.registry.Close()
|
||||
al.GetRegistry().Close()
|
||||
}
|
||||
|
||||
func (al *AgentLoop) RegisterTool(tool tools.Tool) {
|
||||
for _, agentID := range al.registry.ListAgentIDs() {
|
||||
if agent, ok := al.registry.GetAgent(agentID); ok {
|
||||
registry := al.GetRegistry()
|
||||
for _, agentID := range registry.ListAgentIDs() {
|
||||
if agent, ok := registry.GetAgent(agentID); ok {
|
||||
agent.Tools.Register(tool)
|
||||
}
|
||||
}
|
||||
|
|
@ -346,12 +351,123 @@ func (al *AgentLoop) SetChannelManager(cm *channels.Manager) {
|
|||
al.channelManager = cm
|
||||
}
|
||||
|
||||
// ReloadProviderAndConfig atomically swaps the provider and config with proper synchronization.
|
||||
// It uses a context to allow timeout control from the caller.
|
||||
// Returns an error if the reload fails or context is canceled.
|
||||
func (al *AgentLoop) ReloadProviderAndConfig(
|
||||
ctx context.Context,
|
||||
provider providers.LLMProvider,
|
||||
cfg *config.Config,
|
||||
) error {
|
||||
// Validate inputs
|
||||
if provider == nil {
|
||||
return fmt.Errorf("provider cannot be nil")
|
||||
}
|
||||
if cfg == nil {
|
||||
return fmt.Errorf("config cannot be nil")
|
||||
}
|
||||
|
||||
// Create new registry with updated config and provider
|
||||
// Wrap in defer/recover to handle any panics gracefully
|
||||
var registry *AgentRegistry
|
||||
var panicErr error
|
||||
done := make(chan struct{}, 1)
|
||||
|
||||
go func() {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
panicErr = fmt.Errorf("panic during registry creation: %v", r)
|
||||
logger.ErrorCF("agent", "Panic during registry creation",
|
||||
map[string]any{"panic": r})
|
||||
}
|
||||
close(done)
|
||||
}()
|
||||
|
||||
registry = NewAgentRegistry(cfg, provider)
|
||||
}()
|
||||
|
||||
// Wait for completion or context cancellation
|
||||
select {
|
||||
case <-done:
|
||||
if registry == nil {
|
||||
if panicErr != nil {
|
||||
return fmt.Errorf("registry creation failed: %w", panicErr)
|
||||
}
|
||||
return fmt.Errorf("registry creation failed (nil result)")
|
||||
}
|
||||
case <-ctx.Done():
|
||||
return fmt.Errorf("context canceled during registry creation: %w", ctx.Err())
|
||||
}
|
||||
|
||||
// Check context again before proceeding
|
||||
if err := ctx.Err(); err != nil {
|
||||
return fmt.Errorf("context canceled after registry creation: %w", err)
|
||||
}
|
||||
|
||||
// Ensure shared tools are re-registered on the new registry
|
||||
registerSharedTools(cfg, al.bus, registry, provider)
|
||||
|
||||
// Atomically swap the config and registry under write lock
|
||||
// This ensures readers see a consistent pair
|
||||
al.mu.Lock()
|
||||
oldRegistry := al.registry
|
||||
|
||||
// Store new values
|
||||
al.cfg = cfg
|
||||
al.registry = registry
|
||||
|
||||
// Also update fallback chain with new config
|
||||
al.fallback = providers.NewFallbackChain(providers.NewCooldownTracker())
|
||||
|
||||
al.mu.Unlock()
|
||||
|
||||
// Close old provider after releasing the lock
|
||||
// This prevents blocking readers while closing
|
||||
if oldProvider, ok := extractProvider(oldRegistry); ok {
|
||||
if stateful, ok := oldProvider.(providers.StatefulProvider); ok {
|
||||
// Give in-flight requests a moment to complete
|
||||
// Use a reasonable timeout that balances cleanup vs resource usage
|
||||
select {
|
||||
case <-time.After(100 * time.Millisecond):
|
||||
stateful.Close()
|
||||
case <-ctx.Done():
|
||||
// Context canceled, close immediately but log warning
|
||||
logger.WarnCF("agent", "Context canceled during provider cleanup, forcing close",
|
||||
map[string]any{"error": ctx.Err()})
|
||||
stateful.Close()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
logger.InfoCF("agent", "Provider and config reloaded successfully",
|
||||
map[string]any{
|
||||
"model": cfg.Agents.Defaults.GetModelName(),
|
||||
})
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetRegistry returns the current registry (thread-safe)
|
||||
func (al *AgentLoop) GetRegistry() *AgentRegistry {
|
||||
al.mu.RLock()
|
||||
defer al.mu.RUnlock()
|
||||
return al.registry
|
||||
}
|
||||
|
||||
// GetConfig returns the current config (thread-safe)
|
||||
func (al *AgentLoop) GetConfig() *config.Config {
|
||||
al.mu.RLock()
|
||||
defer al.mu.RUnlock()
|
||||
return al.cfg
|
||||
}
|
||||
|
||||
// SetMediaStore injects a MediaStore for media lifecycle management.
|
||||
func (al *AgentLoop) SetMediaStore(s media.MediaStore) {
|
||||
al.mediaStore = s
|
||||
|
||||
// Propagate store to send_file tools in all agents.
|
||||
al.registry.ForEachTool("send_file", func(t tools.Tool) {
|
||||
registry := al.GetRegistry()
|
||||
registry.ForEachTool("send_file", func(t tools.Tool) {
|
||||
if sf, ok := t.(*tools.SendFileTool); ok {
|
||||
sf.SetMediaStore(s)
|
||||
}
|
||||
|
|
@ -540,7 +656,7 @@ func (al *AgentLoop) ProcessHeartbeat(
|
|||
ctx context.Context,
|
||||
content, channel, chatID string,
|
||||
) (string, error) {
|
||||
agent := al.registry.GetDefaultAgent()
|
||||
agent := al.GetRegistry().GetDefaultAgent()
|
||||
if agent == nil {
|
||||
return "", fmt.Errorf("no default agent for heartbeat")
|
||||
}
|
||||
|
|
@ -636,7 +752,8 @@ func (al *AgentLoop) processMessage(ctx context.Context, msg bus.InboundMessage)
|
|||
}
|
||||
|
||||
func (al *AgentLoop) resolveMessageRoute(msg bus.InboundMessage) (routing.ResolvedRoute, *AgentInstance, error) {
|
||||
route := al.registry.ResolveRoute(routing.RouteInput{
|
||||
registry := al.GetRegistry()
|
||||
route := registry.ResolveRoute(routing.RouteInput{
|
||||
Channel: msg.Channel,
|
||||
AccountID: inboundMetadata(msg, metadataKeyAccountID),
|
||||
Peer: extractPeer(msg),
|
||||
|
|
@ -645,9 +762,9 @@ func (al *AgentLoop) resolveMessageRoute(msg bus.InboundMessage) (routing.Resolv
|
|||
TeamID: inboundMetadata(msg, metadataKeyTeamID),
|
||||
})
|
||||
|
||||
agent, ok := al.registry.GetAgent(route.AgentID)
|
||||
agent, ok := registry.GetAgent(route.AgentID)
|
||||
if !ok {
|
||||
agent = al.registry.GetDefaultAgent()
|
||||
agent = registry.GetDefaultAgent()
|
||||
}
|
||||
if agent == nil {
|
||||
return routing.ResolvedRoute{}, nil, fmt.Errorf("no agent available for route (agent_id=%s)", route.AgentID)
|
||||
|
|
@ -709,7 +826,7 @@ func (al *AgentLoop) processSystemMessage(
|
|||
}
|
||||
|
||||
// Use default agent for system messages
|
||||
agent := al.registry.GetDefaultAgent()
|
||||
agent := al.GetRegistry().GetDefaultAgent()
|
||||
if agent == nil {
|
||||
return "", fmt.Errorf("no default agent for system message")
|
||||
}
|
||||
|
|
@ -765,7 +882,8 @@ func (al *AgentLoop) runAgentLoop(
|
|||
)
|
||||
|
||||
// Resolve media:// refs to base64 data URLs (streaming)
|
||||
maxMediaSize := al.cfg.Agents.Defaults.GetMaxMediaSize()
|
||||
cfg := al.GetConfig()
|
||||
maxMediaSize := cfg.Agents.Defaults.GetMaxMediaSize()
|
||||
messages = resolveMediaRefs(messages, al.mediaStore, maxMediaSize)
|
||||
|
||||
// 2. Save user message to session
|
||||
|
|
@ -943,6 +1061,9 @@ func (al *AgentLoop) runLLMIteration(
|
|||
}
|
||||
|
||||
callLLM := func() (*providers.LLMResponse, error) {
|
||||
al.activeRequests.Add(1)
|
||||
defer al.activeRequests.Done()
|
||||
|
||||
if len(activeCandidates) > 1 && al.fallback != nil {
|
||||
fbResult, fbErr := al.fallback.Execute(
|
||||
ctx,
|
||||
|
|
@ -1041,6 +1162,7 @@ func (al *AgentLoop) runLLMIteration(
|
|||
map[string]any{
|
||||
"agent_id": agent.ID,
|
||||
"iteration": iteration,
|
||||
"model": activeModel,
|
||||
"error": err.Error(),
|
||||
})
|
||||
return "", iteration, fmt.Errorf("LLM call failed after retries: %w", err)
|
||||
|
|
@ -1392,7 +1514,8 @@ func (al *AgentLoop) forceCompression(agent *AgentInstance, sessionKey string) {
|
|||
func (al *AgentLoop) GetStartupInfo() map[string]any {
|
||||
info := make(map[string]any)
|
||||
|
||||
agent := al.registry.GetDefaultAgent()
|
||||
registry := al.GetRegistry()
|
||||
agent := registry.GetDefaultAgent()
|
||||
if agent == nil {
|
||||
return info
|
||||
}
|
||||
|
|
@ -1409,8 +1532,8 @@ func (al *AgentLoop) GetStartupInfo() map[string]any {
|
|||
|
||||
// Agents info
|
||||
info["agents"] = map[string]any{
|
||||
"count": len(al.registry.ListAgentIDs()),
|
||||
"ids": al.registry.ListAgentIDs(),
|
||||
"count": len(registry.ListAgentIDs()),
|
||||
"ids": registry.ListAgentIDs(),
|
||||
}
|
||||
|
||||
return info
|
||||
|
|
@ -1598,17 +1721,22 @@ func (al *AgentLoop) retryLLMCall(
|
|||
var err error
|
||||
|
||||
for attempt := 0; attempt < maxRetries; attempt++ {
|
||||
resp, err = agent.Provider.Chat(
|
||||
ctx,
|
||||
[]providers.Message{{Role: "user", Content: prompt}},
|
||||
nil,
|
||||
agent.Model,
|
||||
map[string]any{
|
||||
"max_tokens": agent.MaxTokens,
|
||||
"temperature": llmTemperature,
|
||||
"prompt_cache_key": agent.ID,
|
||||
},
|
||||
)
|
||||
al.activeRequests.Add(1)
|
||||
resp, err = func() (*providers.LLMResponse, error) {
|
||||
defer al.activeRequests.Done()
|
||||
return agent.Provider.Chat(
|
||||
ctx,
|
||||
[]providers.Message{{Role: "user", Content: prompt}},
|
||||
nil,
|
||||
agent.Model,
|
||||
map[string]any{
|
||||
"max_tokens": agent.MaxTokens,
|
||||
"temperature": llmTemperature,
|
||||
"prompt_cache_key": agent.ID,
|
||||
},
|
||||
)
|
||||
}()
|
||||
|
||||
if err == nil && resp != nil && resp.Content != "" {
|
||||
return resp, nil
|
||||
}
|
||||
|
|
@ -1741,9 +1869,11 @@ func (al *AgentLoop) handleCommand(
|
|||
}
|
||||
|
||||
func (al *AgentLoop) buildCommandsRuntime(agent *AgentInstance, opts *processOptions) *commands.Runtime {
|
||||
registry := al.GetRegistry()
|
||||
cfg := al.GetConfig()
|
||||
rt := &commands.Runtime{
|
||||
Config: al.cfg,
|
||||
ListAgentIDs: al.registry.ListAgentIDs,
|
||||
Config: cfg,
|
||||
ListAgentIDs: registry.ListAgentIDs,
|
||||
ListDefinitions: al.cmdRegistry.Definitions,
|
||||
GetEnabledChannels: func() []string {
|
||||
if al.channelManager == nil {
|
||||
|
|
@ -1763,7 +1893,7 @@ func (al *AgentLoop) buildCommandsRuntime(agent *AgentInstance, opts *processOpt
|
|||
}
|
||||
if agent != nil {
|
||||
rt.GetModelInfo = func() (string, string) {
|
||||
return agent.Model, al.cfg.Agents.Defaults.Provider
|
||||
return agent.Model, cfg.Agents.Defaults.Provider
|
||||
}
|
||||
rt.SwitchModel = func(value string) (string, error) {
|
||||
oldModel := agent.Model
|
||||
|
|
@ -1827,3 +1957,16 @@ func extractParentPeer(msg bus.InboundMessage) *routing.RoutePeer {
|
|||
}
|
||||
return &routing.RoutePeer{Kind: parentKind, ID: parentID}
|
||||
}
|
||||
|
||||
// Helper to extract provider from registry for cleanup
|
||||
func extractProvider(registry *AgentRegistry) (providers.LLMProvider, bool) {
|
||||
if registry == nil {
|
||||
return nil, false
|
||||
}
|
||||
// Get any agent to access the provider
|
||||
defaultAgent := registry.GetDefaultAgent()
|
||||
if defaultAgent == nil {
|
||||
return nil, false
|
||||
}
|
||||
return defaultAgent.Provider, true
|
||||
}
|
||||
|
|
|
|||
|
|
@ -222,8 +222,8 @@ type AgentDefaults struct {
|
|||
RestrictToWorkspace bool `json:"restrict_to_workspace" env:"PICOCLAW_AGENTS_DEFAULTS_RESTRICT_TO_WORKSPACE"`
|
||||
AllowReadOutsideWorkspace bool `json:"allow_read_outside_workspace" env:"PICOCLAW_AGENTS_DEFAULTS_ALLOW_READ_OUTSIDE_WORKSPACE"`
|
||||
Provider string `json:"provider" env:"PICOCLAW_AGENTS_DEFAULTS_PROVIDER"`
|
||||
ModelName string `json:"model_name,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_MODEL_NAME"`
|
||||
Model string `json:"model" env:"PICOCLAW_AGENTS_DEFAULTS_MODEL"` // Deprecated: use model_name instead
|
||||
ModelName string `json:"model_name" env:"PICOCLAW_AGENTS_DEFAULTS_MODEL_NAME"`
|
||||
Model string `json:"model,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_MODEL"` // Deprecated: use model_name instead
|
||||
ModelFallbacks []string `json:"model_fallbacks,omitempty"`
|
||||
ImageModel string `json:"image_model,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_IMAGE_MODEL"`
|
||||
ImageModelFallbacks []string `json:"image_model_fallbacks,omitempty"`
|
||||
|
|
@ -528,6 +528,7 @@ type ProvidersConfig struct {
|
|||
Avian ProviderConfig `json:"avian"`
|
||||
Minimax ProviderConfig `json:"minimax"`
|
||||
LongCat ProviderConfig `json:"longcat"`
|
||||
ModelScope ProviderConfig `json:"modelscope"`
|
||||
}
|
||||
|
||||
// IsEmpty checks if all provider configs are empty (no API keys or API bases set)
|
||||
|
|
@ -555,7 +556,8 @@ func (p ProvidersConfig) IsEmpty() bool {
|
|||
p.Mistral.APIKey == "" && p.Mistral.APIBase == "" &&
|
||||
p.Avian.APIKey == "" && p.Avian.APIBase == "" &&
|
||||
p.Minimax.APIKey == "" && p.Minimax.APIBase == "" &&
|
||||
p.LongCat.APIKey == "" && p.LongCat.APIBase == ""
|
||||
p.LongCat.APIKey == "" && p.LongCat.APIBase == "" &&
|
||||
p.ModelScope.APIKey == "" && p.ModelScope.APIBase == ""
|
||||
}
|
||||
|
||||
// MarshalJSON implements custom JSON marshaling for ProvidersConfig
|
||||
|
|
@ -711,6 +713,7 @@ type ExecConfig struct {
|
|||
type SkillsToolsConfig struct {
|
||||
ToolConfig ` envPrefix:"PICOCLAW_TOOLS_SKILLS_"`
|
||||
Registries SkillsRegistriesConfig ` json:"registries"`
|
||||
Github SkillsGithubConfig ` json:"github"`
|
||||
MaxConcurrentSearches int ` json:"max_concurrent_searches" env:"PICOCLAW_TOOLS_SKILLS_MAX_CONCURRENT_SEARCHES"`
|
||||
SearchCache SearchCacheConfig ` json:"search_cache"`
|
||||
}
|
||||
|
|
@ -760,6 +763,11 @@ type SkillsRegistriesConfig struct {
|
|||
ClawHub ClawHubRegistryConfig `json:"clawhub"`
|
||||
}
|
||||
|
||||
type SkillsGithubConfig struct {
|
||||
Token string `json:"token,omitempty" env:"PICOCLAW_TOOLS_SKILLS_GITHUB_AUTH_TOKEN"`
|
||||
Proxy string `json:"proxy,omitempty" env:"PICOCLAW_TOOLS_SKILLS_GITHUB_PROXY"`
|
||||
}
|
||||
|
||||
type ClawHubRegistryConfig struct {
|
||||
Enabled bool `json:"enabled" env:"PICOCLAW_SKILLS_REGISTRIES_CLAWHUB_ENABLED"`
|
||||
BaseURL string `json:"base_url" env:"PICOCLAW_SKILLS_REGISTRIES_CLAWHUB_BASE_URL"`
|
||||
|
|
|
|||
|
|
@ -342,8 +342,8 @@ func TestSaveConfig_IncludesEmptyLegacyModelField(t *testing.T) {
|
|||
t.Fatalf("ReadFile failed: %v", err)
|
||||
}
|
||||
|
||||
if !strings.Contains(string(data), `"model": ""`) {
|
||||
t.Fatalf("saved config should include empty legacy model field, got: %s", string(data))
|
||||
if !strings.Contains(string(data), `"model_name": ""`) {
|
||||
t.Fatalf("saved config should include empty legacy model_name field, got: %s", string(data))
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -369,6 +369,14 @@ func DefaultConfig() *Config {
|
|||
APIKey: "",
|
||||
},
|
||||
|
||||
// ModelScope (魔搭社区) - https://modelscope.cn/my/tokens
|
||||
{
|
||||
ModelName: "modelscope-qwen",
|
||||
Model: "modelscope/Qwen/Qwen3-235B-A22B-Instruct-2507",
|
||||
APIBase: "https://api-inference.modelscope.cn/v1",
|
||||
APIKey: "",
|
||||
},
|
||||
|
||||
// VLLM (local) - http://localhost:8000
|
||||
{
|
||||
ModelName: "local-model",
|
||||
|
|
|
|||
|
|
@ -424,6 +424,23 @@ func ConvertProvidersToModelList(cfg *Config) []ModelConfig {
|
|||
}, true
|
||||
},
|
||||
},
|
||||
{
|
||||
providerNames: []string{"modelscope"},
|
||||
protocol: "modelscope",
|
||||
buildConfig: func(p ProvidersConfig) (ModelConfig, bool) {
|
||||
if p.ModelScope.APIKey == "" && p.ModelScope.APIBase == "" {
|
||||
return ModelConfig{}, false
|
||||
}
|
||||
return ModelConfig{
|
||||
ModelName: "modelscope",
|
||||
Model: "modelscope/Qwen/Qwen3-235B-A22B-Instruct-2507",
|
||||
APIKey: p.ModelScope.APIKey,
|
||||
APIBase: p.ModelScope.APIBase,
|
||||
Proxy: p.ModelScope.Proxy,
|
||||
RequestTimeout: p.ModelScope.RequestTimeout,
|
||||
}, true
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
// Process each provider migration
|
||||
|
|
|
|||
|
|
@ -163,14 +163,15 @@ func TestConvertProvidersToModelList_AllProviders(t *testing.T) {
|
|||
Mistral: ProviderConfig{APIKey: "key18"},
|
||||
Avian: ProviderConfig{APIKey: "key19"},
|
||||
LongCat: ProviderConfig{APIKey: "key-longcat"},
|
||||
ModelScope: ProviderConfig{APIKey: "key-modelscope"},
|
||||
},
|
||||
}
|
||||
|
||||
result := ConvertProvidersToModelList(cfg)
|
||||
|
||||
// All 22 providers should be converted
|
||||
if len(result) != 22 {
|
||||
t.Errorf("len(result) = %d, want 22", len(result))
|
||||
// All 23 providers should be converted
|
||||
if len(result) != 23 {
|
||||
t.Errorf("len(result) = %d, want 23", len(result))
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -59,6 +59,9 @@ func MatchAllowed(sender bus.SenderInfo, allowed string) bool {
|
|||
}
|
||||
}
|
||||
|
||||
// Keep track of explicit username format
|
||||
isAtUsername := strings.HasPrefix(allowed, "@")
|
||||
|
||||
// Strip leading "@" for username matching
|
||||
trimmed := strings.TrimPrefix(allowed, "@")
|
||||
|
||||
|
|
@ -75,11 +78,9 @@ func MatchAllowed(sender bus.SenderInfo, allowed string) bool {
|
|||
return true
|
||||
}
|
||||
|
||||
// Match against Username
|
||||
if sender.Username != "" {
|
||||
if sender.Username == trimmed || sender.Username == allowedUser {
|
||||
return true
|
||||
}
|
||||
// Match against Username only when explicitly requested via "@username"
|
||||
if isAtUsername && sender.Username != "" && sender.Username == trimmed {
|
||||
return true
|
||||
}
|
||||
|
||||
// Match compound sender format against allowed parts
|
||||
|
|
|
|||
|
|
@ -104,6 +104,16 @@ func TestMatchAllowed(t *testing.T) {
|
|||
allowed: "@alice",
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
name: "plain entry does not match username",
|
||||
sender: bus.SenderInfo{
|
||||
Platform: "discord",
|
||||
PlatformID: "999999",
|
||||
Username: "123456",
|
||||
},
|
||||
allowed: "123456",
|
||||
want: false,
|
||||
},
|
||||
{
|
||||
name: "@username does not match",
|
||||
sender: telegramSender,
|
||||
|
|
@ -123,6 +133,16 @@ func TestMatchAllowed(t *testing.T) {
|
|||
allowed: "999|alice",
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
name: "compound matches by ID when username differs",
|
||||
sender: bus.SenderInfo{
|
||||
Platform: "discord",
|
||||
PlatformID: "123456",
|
||||
Username: "not123456",
|
||||
},
|
||||
allowed: "123456|alice",
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
name: "compound does not match",
|
||||
sender: telegramSender,
|
||||
|
|
|
|||
|
|
@ -195,6 +195,10 @@ func DebugC(component string, message string) {
|
|||
logMessage(DEBUG, component, message, nil)
|
||||
}
|
||||
|
||||
func Debugf(message string, ss ...any) {
|
||||
logMessage(DEBUG, "", fmt.Sprintf(message, ss...), nil)
|
||||
}
|
||||
|
||||
func DebugF(message string, fields map[string]any) {
|
||||
logMessage(DEBUG, "", message, fields)
|
||||
}
|
||||
|
|
@ -215,6 +219,10 @@ func InfoF(message string, fields map[string]any) {
|
|||
logMessage(INFO, "", message, fields)
|
||||
}
|
||||
|
||||
func Infof(message string, ss ...any) {
|
||||
logMessage(INFO, "", fmt.Sprintf(message, ss...), nil)
|
||||
}
|
||||
|
||||
func InfoCF(component string, message string, fields map[string]any) {
|
||||
logMessage(INFO, component, message, fields)
|
||||
}
|
||||
|
|
@ -243,6 +251,10 @@ func ErrorC(component string, message string) {
|
|||
logMessage(ERROR, component, message, nil)
|
||||
}
|
||||
|
||||
func Errorf(message string, ss ...any) {
|
||||
logMessage(ERROR, "", fmt.Sprintf(message, ss...), nil)
|
||||
}
|
||||
|
||||
func ErrorF(message string, fields map[string]any) {
|
||||
logMessage(ERROR, "", message, fields)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -123,17 +123,21 @@ func TestLoggerHelperFunctions(t *testing.T) {
|
|||
SetLevel(INFO)
|
||||
|
||||
Debug("This should not log")
|
||||
Debugf("this should not log")
|
||||
Info("This should log")
|
||||
Warn("This should log")
|
||||
Error("This should log")
|
||||
|
||||
InfoC("test", "Component message")
|
||||
InfoF("Fields message", map[string]any{"key": "value"})
|
||||
Infof("test from %v", "Infof")
|
||||
|
||||
WarnC("test", "Warning with component")
|
||||
ErrorF("Error with fields", map[string]any{"error": "test"})
|
||||
Errorf("test from %v", "Errorf")
|
||||
|
||||
SetLevel(DEBUG)
|
||||
DebugC("test", "Debug with component")
|
||||
Debugf("test from %v", "Debugf")
|
||||
WarnF("Warning with fields", map[string]any{"key": "value"})
|
||||
}
|
||||
|
|
|
|||
415
pkg/providers/anthropic_messages/provider.go
Normal file
415
pkg/providers/anthropic_messages/provider.go
Normal file
|
|
@ -0,0 +1,415 @@
|
|||
// PicoClaw - Ultra-lightweight personal AI agent
|
||||
// License: MIT
|
||||
//
|
||||
// Copyright (c) 2026 PicoClaw contributors
|
||||
|
||||
package anthropicmessages
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/providers/protocoltypes"
|
||||
)
|
||||
|
||||
type (
|
||||
ToolCall = protocoltypes.ToolCall
|
||||
FunctionCall = protocoltypes.FunctionCall
|
||||
LLMResponse = protocoltypes.LLMResponse
|
||||
UsageInfo = protocoltypes.UsageInfo
|
||||
Message = protocoltypes.Message
|
||||
ToolDefinition = protocoltypes.ToolDefinition
|
||||
ToolFunctionDefinition = protocoltypes.ToolFunctionDefinition
|
||||
)
|
||||
|
||||
const (
|
||||
defaultAPIVersion = "2023-06-01"
|
||||
defaultBaseURL = "https://api.anthropic.com/v1"
|
||||
defaultRequestTimeout = 120 * time.Second
|
||||
)
|
||||
|
||||
// Provider implements Anthropic Messages API via HTTP (without SDK).
|
||||
// It supports custom endpoints that use Anthropic's native message format.
|
||||
type Provider struct {
|
||||
apiKey string
|
||||
apiBase string
|
||||
httpClient *http.Client
|
||||
}
|
||||
|
||||
// NewProvider creates a new Anthropic Messages API provider.
|
||||
func NewProvider(apiKey, apiBase string) *Provider {
|
||||
return NewProviderWithTimeout(apiKey, apiBase, 0)
|
||||
}
|
||||
|
||||
// NewProviderWithTimeout creates a provider with custom request timeout.
|
||||
func NewProviderWithTimeout(apiKey, apiBase string, timeoutSeconds int) *Provider {
|
||||
baseURL := normalizeBaseURL(apiBase)
|
||||
timeout := defaultRequestTimeout
|
||||
if timeoutSeconds > 0 {
|
||||
timeout = time.Duration(timeoutSeconds) * time.Second
|
||||
}
|
||||
|
||||
return &Provider{
|
||||
apiKey: apiKey,
|
||||
apiBase: baseURL,
|
||||
httpClient: &http.Client{
|
||||
Timeout: timeout,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Chat sends messages to the Anthropic Messages API and returns the response.
|
||||
func (p *Provider) Chat(
|
||||
ctx context.Context,
|
||||
messages []Message,
|
||||
tools []ToolDefinition,
|
||||
model string,
|
||||
options map[string]any,
|
||||
) (*LLMResponse, error) {
|
||||
if p.apiKey == "" {
|
||||
return nil, fmt.Errorf("API key not configured")
|
||||
}
|
||||
|
||||
// Build request body
|
||||
requestBody, err := buildRequestBody(messages, tools, model, options)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("building request body: %w", err)
|
||||
}
|
||||
|
||||
// Serialize to JSON
|
||||
jsonBody, err := json.Marshal(requestBody)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("serializing request body: %w", err)
|
||||
}
|
||||
|
||||
// Build request URL
|
||||
endpointURL, err := url.JoinPath(p.apiBase, "messages")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("building endpoint URL: %w", err)
|
||||
}
|
||||
|
||||
// Create HTTP request
|
||||
req, err := http.NewRequestWithContext(ctx, "POST", endpointURL, bytes.NewReader(jsonBody))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("creating HTTP request: %w", err)
|
||||
}
|
||||
|
||||
// Set headers
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("X-API-Key", p.apiKey) //nolint:canonicalheader // Anthropic API requires exact header name
|
||||
req.Header.Set("Anthropic-Version", defaultAPIVersion)
|
||||
|
||||
// Execute request
|
||||
resp, err := p.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("executing HTTP request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
// Read response body
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("reading response body: %w", err)
|
||||
}
|
||||
|
||||
// Check for HTTP errors with detailed messages
|
||||
switch resp.StatusCode {
|
||||
case http.StatusUnauthorized:
|
||||
return nil, fmt.Errorf("authentication failed (401): check your API key")
|
||||
case http.StatusTooManyRequests:
|
||||
return nil, fmt.Errorf("rate limited (429): %s", string(body))
|
||||
case http.StatusBadRequest:
|
||||
return nil, fmt.Errorf("bad request (400): %s", string(body))
|
||||
case http.StatusNotFound:
|
||||
return nil, fmt.Errorf("endpoint not found (404): %s", string(body))
|
||||
case http.StatusInternalServerError:
|
||||
return nil, fmt.Errorf("internal server error (500): %s", string(body))
|
||||
case http.StatusServiceUnavailable:
|
||||
return nil, fmt.Errorf("service unavailable (503): %s", string(body))
|
||||
default:
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
|
||||
}
|
||||
}
|
||||
|
||||
// Parse response
|
||||
return parseResponseBody(body)
|
||||
}
|
||||
|
||||
// GetDefaultModel returns the default model for this provider.
|
||||
func (p *Provider) GetDefaultModel() string {
|
||||
return "claude-sonnet-4.6"
|
||||
}
|
||||
|
||||
// buildRequestBody converts internal message format to Anthropic Messages API format.
|
||||
func buildRequestBody(
|
||||
messages []Message,
|
||||
tools []ToolDefinition,
|
||||
model string,
|
||||
options map[string]any,
|
||||
) (map[string]any, error) {
|
||||
// max_tokens is required and guaranteed by agent loop
|
||||
maxTokens, ok := asInt(options["max_tokens"])
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("max_tokens is required in options")
|
||||
}
|
||||
|
||||
result := map[string]any{
|
||||
"model": model,
|
||||
"max_tokens": int64(maxTokens),
|
||||
"messages": []any{},
|
||||
}
|
||||
|
||||
// Set temperature from options
|
||||
if temp, ok := asFloat(options["temperature"]); ok {
|
||||
result["temperature"] = temp
|
||||
}
|
||||
|
||||
// Process messages
|
||||
var systemPrompt string
|
||||
var apiMessages []any
|
||||
|
||||
for _, msg := range messages {
|
||||
switch msg.Role {
|
||||
case "system":
|
||||
// Accumulate system messages
|
||||
if systemPrompt != "" {
|
||||
systemPrompt += "\n\n" + msg.Content
|
||||
} else {
|
||||
systemPrompt = msg.Content
|
||||
}
|
||||
|
||||
case "user":
|
||||
if msg.ToolCallID != "" {
|
||||
// Tool result message
|
||||
content := []map[string]any{
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": msg.ToolCallID,
|
||||
"content": msg.Content,
|
||||
},
|
||||
}
|
||||
apiMessages = append(apiMessages, map[string]any{
|
||||
"role": "user",
|
||||
"content": content,
|
||||
})
|
||||
} else {
|
||||
// Regular user message
|
||||
apiMessages = append(apiMessages, map[string]any{
|
||||
"role": "user",
|
||||
"content": msg.Content,
|
||||
})
|
||||
}
|
||||
|
||||
case "assistant":
|
||||
content := []any{}
|
||||
|
||||
// Add text content if present
|
||||
if msg.Content != "" {
|
||||
content = append(content, map[string]any{
|
||||
"type": "text",
|
||||
"text": msg.Content,
|
||||
})
|
||||
}
|
||||
|
||||
// Add tool_use blocks
|
||||
for _, tc := range msg.ToolCalls {
|
||||
toolUse := map[string]any{
|
||||
"type": "tool_use",
|
||||
"id": tc.ID,
|
||||
"name": tc.Name,
|
||||
"input": tc.Arguments,
|
||||
}
|
||||
content = append(content, toolUse)
|
||||
}
|
||||
|
||||
apiMessages = append(apiMessages, map[string]any{
|
||||
"role": "assistant",
|
||||
"content": content,
|
||||
})
|
||||
|
||||
case "tool":
|
||||
// Tool result (alternative format)
|
||||
content := []map[string]any{
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": msg.ToolCallID,
|
||||
"content": msg.Content,
|
||||
},
|
||||
}
|
||||
apiMessages = append(apiMessages, map[string]any{
|
||||
"role": "user",
|
||||
"content": content,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
result["messages"] = apiMessages
|
||||
|
||||
// Set system prompt if present
|
||||
if systemPrompt != "" {
|
||||
result["system"] = systemPrompt
|
||||
}
|
||||
|
||||
// Add tools if present
|
||||
if len(tools) > 0 {
|
||||
result["tools"] = buildTools(tools)
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// buildTools converts tool definitions to Anthropic format.
|
||||
func buildTools(tools []ToolDefinition) []any {
|
||||
result := make([]any, len(tools))
|
||||
for i, tool := range tools {
|
||||
toolDef := map[string]any{
|
||||
"name": tool.Function.Name,
|
||||
"description": tool.Function.Description,
|
||||
"input_schema": tool.Function.Parameters,
|
||||
}
|
||||
result[i] = toolDef
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// parseResponseBody parses Anthropic Messages API response.
|
||||
func parseResponseBody(body []byte) (*LLMResponse, error) {
|
||||
var resp anthropicMessageResponse
|
||||
if err := json.Unmarshal(body, &resp); err != nil {
|
||||
return nil, fmt.Errorf("parsing JSON response: %w", err)
|
||||
}
|
||||
|
||||
// Extract content and tool calls
|
||||
var content strings.Builder
|
||||
toolCalls := make([]ToolCall, 0) // Initialize as empty slice (not nil) for consistent JSON serialization
|
||||
|
||||
for _, block := range resp.Content {
|
||||
switch block.Type {
|
||||
case "text":
|
||||
content.WriteString(block.Text)
|
||||
case "tool_use":
|
||||
argsJSON, _ := json.Marshal(block.Input)
|
||||
toolCalls = append(toolCalls, ToolCall{
|
||||
ID: block.ID,
|
||||
Name: block.Name,
|
||||
Arguments: block.Input,
|
||||
Function: &FunctionCall{
|
||||
Name: block.Name,
|
||||
Arguments: string(argsJSON),
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Map stop_reason
|
||||
finishReason := "stop"
|
||||
switch resp.StopReason {
|
||||
case "tool_use":
|
||||
finishReason = "tool_calls"
|
||||
case "max_tokens":
|
||||
finishReason = "length"
|
||||
case "end_turn":
|
||||
finishReason = "stop"
|
||||
case "stop_sequence":
|
||||
finishReason = "stop"
|
||||
}
|
||||
|
||||
return &LLMResponse{
|
||||
Content: content.String(),
|
||||
ToolCalls: toolCalls,
|
||||
FinishReason: finishReason,
|
||||
Usage: &UsageInfo{
|
||||
PromptTokens: int(resp.Usage.InputTokens),
|
||||
CompletionTokens: int(resp.Usage.OutputTokens),
|
||||
TotalTokens: int(resp.Usage.InputTokens + resp.Usage.OutputTokens),
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// normalizeBaseURL ensures the base URL is properly formatted.
|
||||
// It removes /v1 suffix if present (to avoid duplication) and always appends /v1.
|
||||
// This handles edge cases like "https://api.example.com/v1/proxy" correctly.
|
||||
func normalizeBaseURL(apiBase string) string {
|
||||
base := strings.TrimSpace(apiBase)
|
||||
if base == "" {
|
||||
return defaultBaseURL
|
||||
}
|
||||
|
||||
// Remove trailing slashes
|
||||
base = strings.TrimRight(base, "/")
|
||||
|
||||
// Remove /v1 suffix if present (will be re-added)
|
||||
// This prevents duplication for URLs like "https://api.example.com/v1/proxy"
|
||||
if before, ok := strings.CutSuffix(base, "/v1"); ok {
|
||||
base = before
|
||||
}
|
||||
|
||||
// Ensure we don't have an empty string after cutting
|
||||
if base == "" {
|
||||
return defaultBaseURL
|
||||
}
|
||||
|
||||
// Add /v1 suffix (required by Anthropic Messages API)
|
||||
return base + "/v1"
|
||||
}
|
||||
|
||||
// Helper functions for type conversion
|
||||
|
||||
func asInt(v any) (int, bool) {
|
||||
switch val := v.(type) {
|
||||
case int:
|
||||
return val, true
|
||||
case float64:
|
||||
return int(val), true
|
||||
case int64:
|
||||
return int(val), true
|
||||
default:
|
||||
return 0, false
|
||||
}
|
||||
}
|
||||
|
||||
func asFloat(v any) (float64, bool) {
|
||||
switch val := v.(type) {
|
||||
case float64:
|
||||
return val, true
|
||||
case int:
|
||||
return float64(val), true
|
||||
case int64:
|
||||
return float64(val), true
|
||||
default:
|
||||
return 0, false
|
||||
}
|
||||
}
|
||||
|
||||
// Anthropic API response structures
|
||||
|
||||
type anthropicMessageResponse struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"`
|
||||
Role string `json:"role"`
|
||||
Content []contentBlock `json:"content"`
|
||||
StopReason string `json:"stop_reason"`
|
||||
Model string `json:"model"`
|
||||
Usage usageInfo `json:"usage"`
|
||||
}
|
||||
|
||||
type contentBlock struct {
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text,omitempty"`
|
||||
ID string `json:"id,omitempty"`
|
||||
Name string `json:"name,omitempty"`
|
||||
Input map[string]any `json:"input,omitempty"`
|
||||
}
|
||||
|
||||
type usageInfo struct {
|
||||
InputTokens int64 `json:"input_tokens"`
|
||||
OutputTokens int64 `json:"output_tokens"`
|
||||
}
|
||||
622
pkg/providers/anthropic_messages/provider_test.go
Normal file
622
pkg/providers/anthropic_messages/provider_test.go
Normal file
|
|
@ -0,0 +1,622 @@
|
|||
// PicoClaw - Ultra-lightweight personal AI agent
|
||||
// License: MIT
|
||||
//
|
||||
// Copyright (c) 2026 PicoClaw contributors
|
||||
|
||||
package anthropicmessages
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestBuildRequestBody(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
messages []Message
|
||||
tools []ToolDefinition
|
||||
model string
|
||||
options map[string]any
|
||||
want map[string]any
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "basic user message",
|
||||
messages: []Message{
|
||||
{Role: "user", Content: "Hello, world!"},
|
||||
},
|
||||
model: "test-model",
|
||||
options: map[string]any{
|
||||
"max_tokens": 8192,
|
||||
},
|
||||
want: map[string]any{
|
||||
"model": "test-model",
|
||||
"max_tokens": int64(8192),
|
||||
"messages": []any{
|
||||
map[string]any{
|
||||
"role": "user",
|
||||
"content": "Hello, world!",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "user and assistant messages",
|
||||
messages: []Message{
|
||||
{Role: "user", Content: "What is 2+2?"},
|
||||
{Role: "assistant", Content: "4"},
|
||||
},
|
||||
model: "test-model",
|
||||
options: map[string]any{
|
||||
"max_tokens": 8192,
|
||||
},
|
||||
want: map[string]any{
|
||||
"model": "test-model",
|
||||
"max_tokens": int64(8192),
|
||||
"messages": []any{
|
||||
map[string]any{
|
||||
"role": "user",
|
||||
"content": "What is 2+2?",
|
||||
},
|
||||
map[string]any{
|
||||
"role": "assistant",
|
||||
"content": []any{
|
||||
map[string]any{
|
||||
"type": "text",
|
||||
"text": "4",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "with system message",
|
||||
messages: []Message{
|
||||
{Role: "system", Content: "You are a helpful assistant."},
|
||||
{Role: "user", Content: "Hello"},
|
||||
},
|
||||
model: "test-model",
|
||||
options: map[string]any{
|
||||
"max_tokens": 8192,
|
||||
},
|
||||
want: map[string]any{
|
||||
"model": "test-model",
|
||||
"max_tokens": int64(8192),
|
||||
"system": "You are a helpful assistant.",
|
||||
"messages": []any{
|
||||
map[string]any{
|
||||
"role": "user",
|
||||
"content": "Hello",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "with custom max_tokens and temperature",
|
||||
messages: []Message{
|
||||
{Role: "user", Content: "Test"},
|
||||
},
|
||||
model: "test-model",
|
||||
options: map[string]any{
|
||||
"max_tokens": 2048,
|
||||
"temperature": 0.5,
|
||||
},
|
||||
want: map[string]any{
|
||||
"model": "test-model",
|
||||
"max_tokens": int64(2048),
|
||||
"temperature": 0.5,
|
||||
"messages": []any{
|
||||
map[string]any{
|
||||
"role": "user",
|
||||
"content": "Test",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "missing max_tokens returns error",
|
||||
messages: []Message{
|
||||
{Role: "user", Content: "Test"},
|
||||
},
|
||||
model: "test-model",
|
||||
options: map[string]any{},
|
||||
want: nil,
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "with tools",
|
||||
messages: []Message{
|
||||
{Role: "user", Content: "What's the weather?"},
|
||||
},
|
||||
tools: []ToolDefinition{
|
||||
{
|
||||
Function: ToolFunctionDefinition{
|
||||
Name: "get_weather",
|
||||
Description: "Get current weather",
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"location": map[string]any{
|
||||
"type": "string",
|
||||
"description": "City name",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
model: "test-model",
|
||||
options: map[string]any{
|
||||
"max_tokens": 8192,
|
||||
},
|
||||
want: map[string]any{
|
||||
"model": "test-model",
|
||||
"max_tokens": int64(8192),
|
||||
"messages": []any{
|
||||
map[string]any{
|
||||
"role": "user",
|
||||
"content": "What's the weather?",
|
||||
},
|
||||
},
|
||||
"tools": []any{
|
||||
map[string]any{
|
||||
"name": "get_weather",
|
||||
"description": "Get current weather",
|
||||
"input_schema": map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"location": map[string]any{
|
||||
"type": "string",
|
||||
"description": "City name",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got, err := buildRequestBody(tt.messages, tt.tools, tt.model, tt.options)
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("buildRequestBody() error = %v, wantErr %v", err, tt.wantErr)
|
||||
return
|
||||
}
|
||||
if !reflect.DeepEqual(got, tt.want) {
|
||||
gotJSON, _ := json.MarshalIndent(got, "", " ")
|
||||
wantJSON, _ := json.MarshalIndent(tt.want, "", " ")
|
||||
t.Errorf("buildRequestBody() mismatch:\ngot:\n%s\nwant:\n%s", gotJSON, wantJSON)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseResponseBody(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
body []byte
|
||||
want *LLMResponse
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "basic text response",
|
||||
body: []byte(`{
|
||||
"id": "msg-123",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": "Hello, how can I help?"}
|
||||
],
|
||||
"stop_reason": "end_turn",
|
||||
"model": "test-model",
|
||||
"usage": {
|
||||
"input_tokens": 10,
|
||||
"output_tokens": 5
|
||||
}
|
||||
}`),
|
||||
want: &LLMResponse{
|
||||
Content: "Hello, how can I help?",
|
||||
ToolCalls: []ToolCall{},
|
||||
FinishReason: "stop",
|
||||
Usage: &UsageInfo{
|
||||
PromptTokens: 10,
|
||||
CompletionTokens: 5,
|
||||
TotalTokens: 15,
|
||||
},
|
||||
Reasoning: "",
|
||||
ReasoningDetails: nil,
|
||||
},
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "response with tool use",
|
||||
body: []byte(`{
|
||||
"id": "msg-456",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": "I'll check the weather for you."},
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "toolu-123",
|
||||
"name": "get_weather",
|
||||
"input": {"location": "Tokyo"}
|
||||
}
|
||||
],
|
||||
"stop_reason": "tool_use",
|
||||
"model": "test-model",
|
||||
"usage": {
|
||||
"input_tokens": 20,
|
||||
"output_tokens": 15
|
||||
}
|
||||
}`),
|
||||
want: &LLMResponse{
|
||||
Content: "I'll check the weather for you.",
|
||||
ToolCalls: []ToolCall{
|
||||
{
|
||||
ID: "toolu-123",
|
||||
Name: "get_weather",
|
||||
Arguments: map[string]any{
|
||||
"location": "Tokyo",
|
||||
},
|
||||
Function: &FunctionCall{
|
||||
Name: "get_weather",
|
||||
Arguments: `{"location":"Tokyo"}`,
|
||||
},
|
||||
},
|
||||
},
|
||||
FinishReason: "tool_calls",
|
||||
Usage: &UsageInfo{
|
||||
PromptTokens: 20,
|
||||
CompletionTokens: 15,
|
||||
TotalTokens: 35,
|
||||
},
|
||||
Reasoning: "",
|
||||
ReasoningDetails: nil,
|
||||
},
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "invalid JSON",
|
||||
body: []byte(`invalid json`),
|
||||
want: nil,
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
name: "max_tokens stop reason",
|
||||
body: []byte(`{
|
||||
"id": "msg-789",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": "Partial response"}
|
||||
],
|
||||
"stop_reason": "max_tokens",
|
||||
"model": "test-model",
|
||||
"usage": {
|
||||
"input_tokens": 100,
|
||||
"output_tokens": 4096
|
||||
}
|
||||
}`),
|
||||
want: &LLMResponse{
|
||||
Content: "Partial response",
|
||||
ToolCalls: []ToolCall{},
|
||||
FinishReason: "length",
|
||||
Usage: &UsageInfo{
|
||||
PromptTokens: 100,
|
||||
CompletionTokens: 4096,
|
||||
TotalTokens: 4196,
|
||||
},
|
||||
Reasoning: "",
|
||||
ReasoningDetails: nil,
|
||||
},
|
||||
wantErr: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got, err := parseResponseBody(tt.body)
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("parseResponseBody() error = %v, wantErr %v", err, tt.wantErr)
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// Compare individual fields
|
||||
if got.Content != tt.want.Content {
|
||||
t.Errorf("Content = %q, want %q", got.Content, tt.want.Content)
|
||||
}
|
||||
if got.FinishReason != tt.want.FinishReason {
|
||||
t.Errorf("FinishReason = %q, want %q", got.FinishReason, tt.want.FinishReason)
|
||||
}
|
||||
if got.Usage == nil && tt.want.Usage != nil {
|
||||
t.Errorf("Usage = nil, want non-nil")
|
||||
} else if got.Usage != nil && tt.want.Usage == nil {
|
||||
t.Errorf("Usage = non-nil, want nil")
|
||||
} else if got.Usage != nil && tt.want.Usage != nil {
|
||||
if got.Usage.PromptTokens != tt.want.Usage.PromptTokens {
|
||||
t.Errorf("Usage.PromptTokens = %d, want %d", got.Usage.PromptTokens, tt.want.Usage.PromptTokens)
|
||||
}
|
||||
if got.Usage.CompletionTokens != tt.want.Usage.CompletionTokens {
|
||||
t.Errorf("Usage.CompletionTokens = %d, want %d",
|
||||
got.Usage.CompletionTokens, tt.want.Usage.CompletionTokens)
|
||||
}
|
||||
if got.Usage.TotalTokens != tt.want.Usage.TotalTokens {
|
||||
t.Errorf("Usage.TotalTokens = %d, want %d", got.Usage.TotalTokens, tt.want.Usage.TotalTokens)
|
||||
}
|
||||
}
|
||||
if len(got.ToolCalls) != len(tt.want.ToolCalls) {
|
||||
t.Errorf("ToolCalls length = %d, want %d", len(got.ToolCalls), len(tt.want.ToolCalls))
|
||||
} else {
|
||||
for i := range got.ToolCalls {
|
||||
if got.ToolCalls[i].ID != tt.want.ToolCalls[i].ID {
|
||||
t.Errorf("ToolCalls[%d].ID = %q, want %q",
|
||||
i, got.ToolCalls[i].ID, tt.want.ToolCalls[i].ID)
|
||||
}
|
||||
if got.ToolCalls[i].Name != tt.want.ToolCalls[i].Name {
|
||||
t.Errorf("ToolCalls[%d].Name = %q, want %q",
|
||||
i, got.ToolCalls[i].Name, tt.want.ToolCalls[i].Name)
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeBaseURL(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
apiBase string
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "empty string defaults to official API",
|
||||
apiBase: "",
|
||||
expected: "https://api.anthropic.com/v1",
|
||||
},
|
||||
{
|
||||
name: "URL without /v1 gets it appended",
|
||||
apiBase: "https://api.example.com/anthropic",
|
||||
expected: "https://api.example.com/anthropic/v1",
|
||||
},
|
||||
{
|
||||
name: "URL with /v1 remains unchanged",
|
||||
apiBase: "https://api.example.com/v1",
|
||||
expected: "https://api.example.com/v1",
|
||||
},
|
||||
{
|
||||
name: "URL with trailing slash gets cleaned",
|
||||
apiBase: "https://api.example.com/anthropic/",
|
||||
expected: "https://api.example.com/anthropic/v1",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := normalizeBaseURL(tt.apiBase)
|
||||
if got != tt.expected {
|
||||
t.Errorf("normalizeBaseURL(%q) = %q, want %q", tt.apiBase, got, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewProvider(t *testing.T) {
|
||||
provider := NewProvider("test-key", "https://api.example.com")
|
||||
if provider == nil {
|
||||
t.Fatal("NewProvider() returned nil")
|
||||
}
|
||||
if provider.apiKey != "test-key" {
|
||||
t.Errorf("provider.apiKey = %q, want %q", provider.apiKey, "test-key")
|
||||
}
|
||||
if provider.apiBase != "https://api.example.com/v1" {
|
||||
t.Errorf("provider.apiBase = %q, want %q", provider.apiBase, "https://api.example.com/v1")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetDefaultModel(t *testing.T) {
|
||||
provider := NewProvider("test-key", "")
|
||||
got := provider.GetDefaultModel()
|
||||
expected := "claude-sonnet-4.6"
|
||||
if got != expected {
|
||||
t.Errorf("GetDefaultModel() = %q, want %q", got, expected)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildRequestBodyEdgeCases tests edge cases for buildRequestBody.
|
||||
func TestBuildRequestBodyEdgeCases(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
messages []Message
|
||||
tools []ToolDefinition
|
||||
model string
|
||||
options map[string]any
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "empty message list",
|
||||
messages: []Message{},
|
||||
model: "test-model",
|
||||
options: map[string]any{
|
||||
"max_tokens": 8192,
|
||||
},
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "very long system message",
|
||||
messages: []Message{
|
||||
{Role: "system", Content: strings.Repeat("This is a very long system prompt. ", 1000)},
|
||||
{Role: "user", Content: "Hello"},
|
||||
},
|
||||
model: "test-model",
|
||||
options: map[string]any{
|
||||
"max_tokens": 8192,
|
||||
},
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "multiple consecutive system messages",
|
||||
messages: []Message{
|
||||
{Role: "system", Content: "First system message"},
|
||||
{Role: "system", Content: "Second system message"},
|
||||
{Role: "system", Content: "Third system message"},
|
||||
{Role: "user", Content: "Hello"},
|
||||
},
|
||||
model: "test-model",
|
||||
options: map[string]any{
|
||||
"max_tokens": 8192,
|
||||
},
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "tool result without tool call",
|
||||
messages: []Message{
|
||||
{Role: "user", Content: "Use a tool"},
|
||||
{Role: "assistant", Content: "", ToolCalls: []ToolCall{
|
||||
{ID: "tool-1", Name: "test_tool", Arguments: map[string]any{"arg": "value"}},
|
||||
}},
|
||||
{Role: "user", ToolCallID: "tool-1", Content: "Tool result"},
|
||||
},
|
||||
model: "test-model",
|
||||
options: map[string]any{
|
||||
"max_tokens": 8192,
|
||||
},
|
||||
wantErr: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got, err := buildRequestBody(tt.messages, tt.tools, tt.model, tt.options)
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("buildRequestBody() error = %v, wantErr %v", err, tt.wantErr)
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// Verify basic structure
|
||||
if got == nil {
|
||||
t.Error("buildRequestBody() returned nil")
|
||||
return
|
||||
}
|
||||
if got["model"] != tt.model {
|
||||
t.Errorf("model = %v, want %v", got["model"], tt.model)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestParseResponseBodyEdgeCases tests edge cases for parseResponseBody.
|
||||
func TestParseResponseBodyEdgeCases(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
body []byte
|
||||
wantErr bool
|
||||
check func(*testing.T, *LLMResponse)
|
||||
}{
|
||||
{
|
||||
name: "empty content blocks",
|
||||
body: []byte(`{
|
||||
"id": "msg-empty",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [],
|
||||
"stop_reason": "end_turn",
|
||||
"model": "test-model",
|
||||
"usage": {"input_tokens": 5, "output_tokens": 0}
|
||||
}`),
|
||||
wantErr: false,
|
||||
check: func(t *testing.T, resp *LLMResponse) {
|
||||
if resp.Content != "" {
|
||||
t.Errorf("Content = %q, want empty string", resp.Content)
|
||||
}
|
||||
if len(resp.ToolCalls) != 0 {
|
||||
t.Errorf("ToolCalls length = %d, want 0", len(resp.ToolCalls))
|
||||
}
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "multiple tool use blocks",
|
||||
body: []byte(`{
|
||||
"id": "msg-multi",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "tool_use", "id": "tool-1", "name": "func1", "input": {"arg": "val1"}},
|
||||
{"type": "tool_use", "id": "tool-2", "name": "func2", "input": {"arg": "val2"}}
|
||||
],
|
||||
"stop_reason": "tool_use",
|
||||
"model": "test-model",
|
||||
"usage": {"input_tokens": 10, "output_tokens": 20}
|
||||
}`),
|
||||
wantErr: false,
|
||||
check: func(t *testing.T, resp *LLMResponse) {
|
||||
if len(resp.ToolCalls) != 2 {
|
||||
t.Errorf("ToolCalls length = %d, want 2", len(resp.ToolCalls))
|
||||
}
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "malformed JSON response",
|
||||
body: []byte(`{invalid json`),
|
||||
wantErr: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got, err := parseResponseBody(tt.body)
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("parseResponseBody() error = %v, wantErr %v", err, tt.wantErr)
|
||||
return
|
||||
}
|
||||
if tt.check != nil && err == nil {
|
||||
tt.check(t, got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestProviderChatErrors tests error handling in Chat.
|
||||
// Note: apiBase check removed as it's dead code - normalizeBaseURL() always provides a default.
|
||||
func TestProviderChatErrors(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
apiKey string
|
||||
messages []Message
|
||||
wantErrMsg string
|
||||
}{
|
||||
{
|
||||
name: "missing API key",
|
||||
apiKey: "",
|
||||
messages: []Message{{Role: "user", Content: "Test"}},
|
||||
wantErrMsg: "API key not configured",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
// Create provider using constructor to ensure proper initialization
|
||||
provider := NewProvider(tt.apiKey, "https://api.example.com")
|
||||
|
||||
_, err := provider.Chat(context.Background(), tt.messages, nil, "test-model", nil)
|
||||
if err == nil {
|
||||
t.Fatal("Chat() expected error, got nil")
|
||||
}
|
||||
if err.Error() != tt.wantErrMsg {
|
||||
t.Errorf("Chat() error = %q, want %q", err.Error(), tt.wantErrMsg)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
@ -11,6 +11,7 @@ import (
|
|||
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
"github.com/sipeed/picoclaw/pkg/providers/azure"
|
||||
anthropicmessages "github.com/sipeed/picoclaw/pkg/providers/anthropic_messages"
|
||||
)
|
||||
|
||||
// createClaudeAuthProvider creates a Claude provider using OAuth credentials from auth store.
|
||||
|
|
@ -54,7 +55,8 @@ func ExtractProtocol(model string) (protocol, modelID string) {
|
|||
|
||||
// CreateProviderFromConfig creates a provider based on the ModelConfig.
|
||||
// It uses the protocol prefix in the Model field to determine which provider to create.
|
||||
// Supported protocols: openai, litellm, anthropic, antigravity, claude-cli, codex-cli, github-copilot
|
||||
// Supported protocols: openai, litellm, anthropic, anthropic-messages, antigravity,
|
||||
// claude-cli, codex-cli, github-copilot
|
||||
// Returns the provider, the model ID (without protocol prefix), and any error.
|
||||
func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, error) {
|
||||
if cfg == nil {
|
||||
|
|
@ -112,7 +114,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err
|
|||
case "litellm", "openrouter", "groq", "zhipu", "gemini", "nvidia",
|
||||
"ollama", "moonshot", "shengsuanyun", "deepseek", "cerebras",
|
||||
"vivgrid", "volcengine", "vllm", "qwen", "mistral", "avian",
|
||||
"minimax", "longcat":
|
||||
"minimax", "longcat", "modelscope":
|
||||
// All other OpenAI-compatible HTTP providers
|
||||
if cfg.APIKey == "" && cfg.APIBase == "" {
|
||||
return nil, "", fmt.Errorf("api_key or api_base is required for HTTP-based protocol %q", protocol)
|
||||
|
|
@ -154,6 +156,21 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err
|
|||
cfg.RequestTimeout,
|
||||
), modelID, nil
|
||||
|
||||
case "anthropic-messages":
|
||||
// Anthropic Messages API with native format (HTTP-based, no SDK)
|
||||
apiBase := cfg.APIBase
|
||||
if apiBase == "" {
|
||||
apiBase = "https://api.anthropic.com/v1"
|
||||
}
|
||||
if cfg.APIKey == "" {
|
||||
return nil, "", fmt.Errorf("api_key is required for anthropic-messages protocol (model: %s)", cfg.Model)
|
||||
}
|
||||
return anthropicmessages.NewProviderWithTimeout(
|
||||
cfg.APIKey,
|
||||
apiBase,
|
||||
cfg.RequestTimeout,
|
||||
), modelID, nil
|
||||
|
||||
case "antigravity":
|
||||
return NewAntigravityProvider(), modelID, nil
|
||||
|
||||
|
|
@ -234,6 +251,8 @@ func getDefaultAPIBase(protocol string) string {
|
|||
return "https://api.minimaxi.com/v1"
|
||||
case "longcat":
|
||||
return "https://api.longcat.chat/openai"
|
||||
case "modelscope":
|
||||
return "https://api-inference.modelscope.cn/v1"
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
|
|
|
|||
|
|
@ -120,6 +120,7 @@ func TestCreateProviderFromConfig_DefaultAPIBase(t *testing.T) {
|
|||
{"deepseek", "deepseek"},
|
||||
{"ollama", "ollama"},
|
||||
{"longcat", "longcat"},
|
||||
{"modelscope", "modelscope"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
|
|
@ -192,6 +193,35 @@ func TestCreateProviderFromConfig_LongCat(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestCreateProviderFromConfig_ModelScope(t *testing.T) {
|
||||
cfg := &config.ModelConfig{
|
||||
ModelName: "test-modelscope",
|
||||
Model: "modelscope/Qwen/Qwen3-235B-A22B-Instruct-2507",
|
||||
APIKey: "test-key",
|
||||
APIBase: "https://api-inference.modelscope.cn/v1",
|
||||
}
|
||||
|
||||
provider, modelID, err := CreateProviderFromConfig(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("CreateProviderFromConfig() error = %v", err)
|
||||
}
|
||||
if provider == nil {
|
||||
t.Fatal("CreateProviderFromConfig() returned nil provider")
|
||||
}
|
||||
if modelID != "Qwen/Qwen3-235B-A22B-Instruct-2507" {
|
||||
t.Errorf("modelID = %q, want %q", modelID, "Qwen/Qwen3-235B-A22B-Instruct-2507")
|
||||
}
|
||||
if _, ok := provider.(*HTTPProvider); !ok {
|
||||
t.Fatalf("expected *HTTPProvider, got %T", provider)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetDefaultAPIBase_ModelScope(t *testing.T) {
|
||||
if got := getDefaultAPIBase("modelscope"); got != "https://api-inference.modelscope.cn/v1" {
|
||||
t.Fatalf("getDefaultAPIBase(%q) = %q, want %q", "modelscope", got, "https://api-inference.modelscope.cn/v1")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateProviderFromConfig_Anthropic(t *testing.T) {
|
||||
cfg := &config.ModelConfig{
|
||||
ModelName: "test-anthropic",
|
||||
|
|
|
|||
|
|
@ -2,80 +2,289 @@ package skills
|
|||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"path"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/fileutil"
|
||||
"github.com/sipeed/picoclaw/pkg/utils"
|
||||
)
|
||||
|
||||
type SkillInstaller struct {
|
||||
workspace string
|
||||
// GitHubContent represents a file or directory in GitHub API response
|
||||
type GitHubContent struct {
|
||||
Name string `json:"name"`
|
||||
Path string `json:"path"`
|
||||
Type string `json:"type"` // "file" or "dir"
|
||||
DownloadURL string `json:"download_url"`
|
||||
URL string `json:"url"` // API URL for subdirectories
|
||||
}
|
||||
|
||||
func NewSkillInstaller(workspace string) *SkillInstaller {
|
||||
return &SkillInstaller{
|
||||
workspace: workspace,
|
||||
// GitHubRef represents a parsed GitHub reference
|
||||
type GitHubRef struct {
|
||||
Owner string // Repository owner
|
||||
RepoName string // Repository name
|
||||
Ref string // Git reference (branch, tag, or commit)
|
||||
SubPath string // Path within the repository
|
||||
}
|
||||
|
||||
type SkillInstaller struct {
|
||||
workspace string
|
||||
client *http.Client
|
||||
githubToken string
|
||||
proxy string
|
||||
}
|
||||
|
||||
// NewSkillInstaller creates a new skill installer.
|
||||
// proxy is an optional HTTP/HTTPS/SOCKS5 proxy URL for downloading skills.
|
||||
func NewSkillInstaller(workspace, githubToken, proxy string) (*SkillInstaller, error) {
|
||||
client, err := utils.CreateHTTPClient(proxy, 15*time.Second)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create HTTP client: %w", err)
|
||||
}
|
||||
|
||||
return &SkillInstaller{
|
||||
workspace: workspace,
|
||||
client: client,
|
||||
githubToken: githubToken,
|
||||
proxy: proxy,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// parseGitHubRef parses a GitHub reference.
|
||||
// Supports: "owner/repo", "owner/repo/path", or full URL like "https://github.com/owner/repo/tree/ref/path"
|
||||
func parseGitHubRef(repo string) (GitHubRef, error) {
|
||||
repo = strings.TrimSpace(repo)
|
||||
|
||||
// Handle full URL
|
||||
if strings.HasPrefix(repo, "http://") || strings.HasPrefix(repo, "https://") {
|
||||
u, err := url.Parse(repo)
|
||||
if err != nil {
|
||||
return GitHubRef{}, fmt.Errorf("invalid URL: %w", err)
|
||||
}
|
||||
parts := strings.Split(strings.Trim(u.Path, "/"), "/")
|
||||
if len(parts) < 2 {
|
||||
return GitHubRef{}, fmt.Errorf("invalid GitHub URL")
|
||||
}
|
||||
ref := GitHubRef{
|
||||
Owner: parts[0],
|
||||
RepoName: parts[1],
|
||||
Ref: "main",
|
||||
}
|
||||
// Look for /tree/ or /blob/ in the path
|
||||
for i := 2; i < len(parts); i++ {
|
||||
if parts[i] == "tree" || parts[i] == "blob" {
|
||||
if i+1 < len(parts) {
|
||||
ref.Ref = parts[i+1]
|
||||
ref.SubPath = strings.Join(parts[i+2:], "/")
|
||||
}
|
||||
break
|
||||
}
|
||||
}
|
||||
return ref, nil
|
||||
}
|
||||
|
||||
// Handle shorthand format
|
||||
parts := strings.Split(strings.Trim(repo, "/"), "/")
|
||||
if len(parts) < 2 {
|
||||
return GitHubRef{}, fmt.Errorf("invalid format %q: expected 'owner/repo'", repo)
|
||||
}
|
||||
ref := GitHubRef{
|
||||
Owner: parts[0],
|
||||
RepoName: parts[1],
|
||||
Ref: "main",
|
||||
}
|
||||
if len(parts) > 2 {
|
||||
ref.SubPath = strings.Join(parts[2:], "/")
|
||||
}
|
||||
return ref, nil
|
||||
}
|
||||
|
||||
func (si *SkillInstaller) InstallFromGitHub(ctx context.Context, repo string) error {
|
||||
skillDir := filepath.Join(si.workspace, "skills", filepath.Base(repo))
|
||||
|
||||
if _, err := os.Stat(skillDir); err == nil {
|
||||
return fmt.Errorf("skill '%s' already exists", filepath.Base(repo))
|
||||
ref, err := parseGitHubRef(repo)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
url := fmt.Sprintf("https://raw.githubusercontent.com/%s/main/SKILL.md", repo)
|
||||
skillName := ref.RepoName
|
||||
if ref.SubPath != "" {
|
||||
skillName = filepath.Base(ref.SubPath)
|
||||
}
|
||||
skillDirectory := filepath.Join(si.workspace, "skills", skillName)
|
||||
|
||||
if _, err := os.Stat(skillDirectory); err == nil {
|
||||
return fmt.Errorf("skill '%s' already exists", skillName)
|
||||
}
|
||||
|
||||
// Build GitHub API URL
|
||||
apiPath := path.Join(ref.Owner, ref.RepoName, "contents")
|
||||
if ref.SubPath != "" {
|
||||
apiPath = path.Join(apiPath, ref.SubPath)
|
||||
}
|
||||
apiURL := fmt.Sprintf("https://api.github.com/repos/%s?ref=%s", apiPath, ref.Ref)
|
||||
|
||||
if err := si.getGithubDirAllFiles(ctx, apiURL, skillDirectory, true); err != nil {
|
||||
// Fallback to raw download
|
||||
return si.downloadRaw(ctx, ref.Owner, ref.RepoName, ref.Ref, ref.SubPath, skillDirectory)
|
||||
}
|
||||
|
||||
if _, err := os.Stat(filepath.Join(skillDirectory, "SKILL.md")); err != nil {
|
||||
return fmt.Errorf("SKILL.md not found in repository")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// downloadDir recursively downloads a directory from GitHub API
|
||||
// isRoot: true if this is the skill root directory (only download SKILL.md at root)
|
||||
func (si *SkillInstaller) getGithubDirAllFiles(ctx context.Context, apiURL, localDir string, isRoot bool) error {
|
||||
req, err := http.NewRequestWithContext(ctx, "GET", apiURL, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if si.githubToken != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+si.githubToken)
|
||||
}
|
||||
|
||||
resp, err := utils.DoRequestWithRetry(si.client, req)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != 200 {
|
||||
return fmt.Errorf("HTTP %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
var items []GitHubContent
|
||||
if err := json.NewDecoder(resp.Body).Decode(&items); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
for _, item := range items {
|
||||
localPath := filepath.Join(localDir, item.Name)
|
||||
|
||||
switch item.Type {
|
||||
case "file":
|
||||
if !shouldDownload(item.Name, isRoot) {
|
||||
continue
|
||||
}
|
||||
if err := si.downloadFile(ctx, item.DownloadURL, localPath); err != nil {
|
||||
return fmt.Errorf("download %s: %w", item.Name, err)
|
||||
}
|
||||
case "dir":
|
||||
if !isSkillDirectory(item.Name) {
|
||||
continue
|
||||
}
|
||||
if err := si.getGithubDirAllFiles(ctx, item.URL, localPath, false); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// downloadRaw is a fallback that downloads just SKILL.md from raw.githubusercontent.com
|
||||
func (si *SkillInstaller) downloadRaw(ctx context.Context, owner, repo, ref, subPath, localDir string) error {
|
||||
urlPath := path.Join(owner, repo, ref)
|
||||
if subPath != "" {
|
||||
urlPath = path.Join(urlPath, subPath)
|
||||
}
|
||||
url := fmt.Sprintf("https://raw.githubusercontent.com/%s/SKILL.md", urlPath)
|
||||
|
||||
client := &http.Client{Timeout: 15 * time.Second}
|
||||
req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
|
||||
resp, err := utils.DoRequestWithRetry(client, req)
|
||||
// Use chunked download to temporary file.
|
||||
tmpPath, err := utils.DownloadToFile(ctx, si.client, req, 0)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to fetch skill: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
defer os.Remove(tmpPath)
|
||||
|
||||
if resp.StatusCode != 200 {
|
||||
return fmt.Errorf("failed to fetch skill: HTTP %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to read response: %w", err)
|
||||
}
|
||||
|
||||
if err := os.MkdirAll(skillDir, 0o755); err != nil {
|
||||
if err := os.MkdirAll(localDir, 0o755); err != nil {
|
||||
return fmt.Errorf("failed to create skill directory: %w", err)
|
||||
}
|
||||
|
||||
skillPath := filepath.Join(skillDir, "SKILL.md")
|
||||
localPath := filepath.Join(localDir, "SKILL.md")
|
||||
|
||||
// Use unified atomic write utility with explicit sync for flash storage reliability.
|
||||
if err := fileutil.WriteFileAtomic(skillPath, body, 0o600); err != nil {
|
||||
// Atomic move from temp to final location.
|
||||
if err := os.Rename(tmpPath, localPath); err != nil {
|
||||
return fmt.Errorf("failed to write skill file: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
return os.Chmod(localPath, 0o600)
|
||||
}
|
||||
|
||||
func (si *SkillInstaller) downloadFile(ctx context.Context, url, localPath string) error {
|
||||
req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Use chunked download to temporary file, then move atomically to target.
|
||||
tmpPath, err := utils.DownloadToFile(ctx, si.client, req, 0)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer os.Remove(tmpPath)
|
||||
|
||||
if err := os.MkdirAll(filepath.Dir(localPath), 0o755); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Atomic move from temp to final location.
|
||||
if err := os.Rename(tmpPath, localPath); err != nil {
|
||||
return fmt.Errorf("failed to move downloaded file: %w", err)
|
||||
}
|
||||
|
||||
return os.Chmod(localPath, 0o600)
|
||||
}
|
||||
|
||||
// shouldDownload determines if a file should be downloaded
|
||||
// root: true if we're at the skill root directory
|
||||
func shouldDownload(name string, root bool) bool {
|
||||
if root {
|
||||
return name == "SKILL.md"
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// isSkillDir checks if a directory is a standard skill resource directory
|
||||
func isSkillDirectory(name string) bool {
|
||||
switch name {
|
||||
case "scripts", "references", "assets", "templates", "docs":
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (si *SkillInstaller) Uninstall(skillName string) error {
|
||||
skillDir := filepath.Join(si.workspace, "skills", skillName)
|
||||
parts := strings.Split(skillName, "/")
|
||||
var finalSkillName string
|
||||
for i := len(parts) - 1; i >= 0; i-- {
|
||||
if parts[i] != "" {
|
||||
finalSkillName = parts[i]
|
||||
break
|
||||
}
|
||||
}
|
||||
if finalSkillName == "" {
|
||||
finalSkillName = skillName
|
||||
}
|
||||
|
||||
skillDir := filepath.Join(si.workspace, "skills", finalSkillName)
|
||||
|
||||
if _, err := os.Stat(skillDir); os.IsNotExist(err) {
|
||||
return fmt.Errorf("skill '%s' not found", skillName)
|
||||
return fmt.Errorf("skill '%s' not found (processed as '%s')", skillName, finalSkillName)
|
||||
}
|
||||
|
||||
if err := os.RemoveAll(skillDir); err != nil {
|
||||
return fmt.Errorf("failed to remove skill: %w", err)
|
||||
return fmt.Errorf("failed to remove skill '%s': %w", finalSkillName, err)
|
||||
}
|
||||
|
||||
return nil
|
||||
|
|
|
|||
665
pkg/skills/installer_test.go
Normal file
665
pkg/skills/installer_test.go
Normal file
|
|
@ -0,0 +1,665 @@
|
|||
package skills
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestParseGitHubRef(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
repo string
|
||||
wantOwner string
|
||||
wantRepoName string
|
||||
wantRef string
|
||||
wantSubPath string
|
||||
wantErr bool
|
||||
wantErrContain string
|
||||
}{
|
||||
{
|
||||
name: "simple owner/repo",
|
||||
repo: "sipeed/picoclaw",
|
||||
wantOwner: "sipeed",
|
||||
wantRepoName: "picoclaw",
|
||||
wantRef: "main",
|
||||
wantSubPath: "",
|
||||
},
|
||||
{
|
||||
name: "owner/repo with subpath",
|
||||
repo: "sipeed/picoclaw/skills/test",
|
||||
wantOwner: "sipeed",
|
||||
wantRepoName: "picoclaw",
|
||||
wantRef: "main",
|
||||
wantSubPath: "skills/test",
|
||||
},
|
||||
{
|
||||
name: "full URL with tree",
|
||||
repo: "https://github.com/sipeed/picoclaw/tree/dev/skills/test",
|
||||
wantOwner: "sipeed",
|
||||
wantRepoName: "picoclaw",
|
||||
wantRef: "dev",
|
||||
wantSubPath: "skills/test",
|
||||
},
|
||||
{
|
||||
name: "full URL with blob",
|
||||
repo: "https://github.com/sipeed/picoclaw/blob/main/README.md",
|
||||
wantOwner: "sipeed",
|
||||
wantRepoName: "picoclaw",
|
||||
wantRef: "main",
|
||||
wantSubPath: "README.md",
|
||||
},
|
||||
{
|
||||
name: "full URL without ref",
|
||||
repo: "https://github.com/sipeed/picoclaw",
|
||||
wantOwner: "sipeed",
|
||||
wantRepoName: "picoclaw",
|
||||
wantRef: "main",
|
||||
wantSubPath: "",
|
||||
},
|
||||
{
|
||||
name: "invalid format - single part",
|
||||
repo: "sipeed",
|
||||
wantErr: true,
|
||||
wantErrContain: "expected 'owner/repo'",
|
||||
},
|
||||
{
|
||||
name: "invalid URL",
|
||||
repo: "http://[invalid",
|
||||
wantErr: true,
|
||||
wantErrContain: "invalid URL",
|
||||
},
|
||||
{
|
||||
name: "invalid GitHub URL - only one path part",
|
||||
repo: "https://github.com/sipeed",
|
||||
wantErr: true,
|
||||
wantErrContain: "invalid GitHub URL",
|
||||
},
|
||||
{
|
||||
name: "with whitespace",
|
||||
repo: " sipeed/picoclaw ",
|
||||
wantOwner: "sipeed",
|
||||
wantRepoName: "picoclaw",
|
||||
wantRef: "main",
|
||||
wantSubPath: "",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
ref, err := parseGitHubRef(tt.repo)
|
||||
|
||||
if tt.wantErr {
|
||||
if err == nil {
|
||||
t.Errorf("parseGitHubRef() error = nil, wantErr = true")
|
||||
return
|
||||
}
|
||||
if tt.wantErrContain != "" && !strings.Contains(err.Error(), tt.wantErrContain) {
|
||||
t.Errorf("parseGitHubRef() error = %v, want error containing %v", err, tt.wantErrContain)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
t.Errorf("parseGitHubRef() unexpected error = %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
if ref.Owner != tt.wantOwner {
|
||||
t.Errorf("parseGitHubRef() owner = %v, want %v", ref.Owner, tt.wantOwner)
|
||||
}
|
||||
if ref.RepoName != tt.wantRepoName {
|
||||
t.Errorf("parseGitHubRef() repoName = %v, want %v", ref.RepoName, tt.wantRepoName)
|
||||
}
|
||||
if ref.Ref != tt.wantRef {
|
||||
t.Errorf("parseGitHubRef() ref = %v, want %v", ref.Ref, tt.wantRef)
|
||||
}
|
||||
if ref.SubPath != tt.wantSubPath {
|
||||
t.Errorf("parseGitHubRef() subPath = %v, want %v", ref.SubPath, tt.wantSubPath)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestShouldDownload(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
file string
|
||||
root bool
|
||||
want bool
|
||||
}{
|
||||
{"SKILL.md at root", "SKILL.md", true, true},
|
||||
{"other file at root", "README.md", true, false},
|
||||
{"script at root", "script.py", true, false},
|
||||
{"SKILL.md not at root", "SKILL.md", false, true},
|
||||
{"any file not at root", "any.txt", false, true},
|
||||
{"script not at root", "script.py", false, true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := shouldDownload(tt.file, tt.root)
|
||||
if got != tt.want {
|
||||
t.Errorf("shouldDownload(%q, %v) = %v, want %v", tt.file, tt.root, got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsSkillDirectory(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
dir string
|
||||
want bool
|
||||
}{
|
||||
{"scripts dir", "scripts", true},
|
||||
{"references dir", "references", true},
|
||||
{"assets dir", "assets", true},
|
||||
{"templates dir", "templates", true},
|
||||
{"docs dir", "docs", true},
|
||||
{"other dir", "other", false},
|
||||
{"src dir", "src", false},
|
||||
{"empty string", "", false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := isSkillDirectory(tt.dir)
|
||||
if got != tt.want {
|
||||
t.Errorf("isSkillDirectory(%q) = %v, want %v", tt.dir, got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewSkillInstaller(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
installer, err := NewSkillInstaller(tmpDir, "test-token", "")
|
||||
if err != nil {
|
||||
t.Fatalf("NewSkillInstaller() error = %v", err)
|
||||
}
|
||||
|
||||
if installer == nil {
|
||||
t.Fatal("NewSkillInstaller() returned nil")
|
||||
}
|
||||
|
||||
if installer.workspace != tmpDir {
|
||||
t.Errorf("workspace = %v, want %v", installer.workspace, tmpDir)
|
||||
}
|
||||
|
||||
if installer.githubToken != "test-token" {
|
||||
t.Errorf("githubToken = %v, want 'test-token'", installer.githubToken)
|
||||
}
|
||||
|
||||
if installer.proxy != "" {
|
||||
t.Errorf("proxy = %v, want empty", installer.proxy)
|
||||
}
|
||||
|
||||
if installer.client == nil {
|
||||
t.Error("client is nil")
|
||||
} else if installer.client.Timeout != 15*time.Second {
|
||||
t.Errorf("client.Timeout = %v, want 15s", installer.client.Timeout)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewSkillInstaller_WithProxy(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
installer, err := NewSkillInstaller(tmpDir, "test-token", "http://127.0.0.1:7890")
|
||||
if err != nil {
|
||||
t.Fatalf("NewSkillInstaller() error = %v", err)
|
||||
}
|
||||
|
||||
if installer.proxy != "http://127.0.0.1:7890" {
|
||||
t.Errorf("proxy = %v, want 'http://127.0.0.1:7890'", installer.proxy)
|
||||
}
|
||||
|
||||
if installer.client == nil {
|
||||
t.Fatal("client is nil")
|
||||
}
|
||||
|
||||
// Verify the transport has proxy configured
|
||||
transport, ok := installer.client.Transport.(*http.Transport)
|
||||
if !ok {
|
||||
t.Fatal("client.Transport is not *http.Transport")
|
||||
}
|
||||
|
||||
if transport.Proxy == nil {
|
||||
t.Error("transport.Proxy is nil, expected non-nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewSkillInstaller_InvalidProxy(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
installer, err := NewSkillInstaller(tmpDir, "test-token", "://invalid-proxy")
|
||||
if err == nil {
|
||||
t.Error("NewSkillInstaller() expected error for invalid proxy, got nil")
|
||||
}
|
||||
if installer != nil {
|
||||
t.Error("expected nil installer on error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSkillInstaller_DownloadFile(t *testing.T) {
|
||||
// Create a test server that serves files
|
||||
content := "test file content for skill download"
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodGet {
|
||||
t.Errorf("expected GET, got %s", r.Method)
|
||||
}
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte(content))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
tmpDir := t.TempDir()
|
||||
installer, err := NewSkillInstaller(tmpDir, "", "")
|
||||
if err != nil {
|
||||
t.Fatalf("NewSkillInstaller() error = %v", err)
|
||||
}
|
||||
|
||||
t.Run("successful download", func(t *testing.T) {
|
||||
localPath := filepath.Join(tmpDir, "test-skill", "SKILL.md")
|
||||
err := installer.downloadFile(context.Background(), server.URL, localPath)
|
||||
if err != nil {
|
||||
t.Errorf("downloadFile() error = %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
// Verify file was downloaded
|
||||
data, err := os.ReadFile(localPath)
|
||||
if err != nil {
|
||||
t.Errorf("failed to read downloaded file: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
if string(data) != content {
|
||||
t.Errorf("downloaded content = %q, want %q", string(data), content)
|
||||
}
|
||||
|
||||
// Check file permissions
|
||||
info, err := os.Stat(localPath)
|
||||
if err != nil {
|
||||
t.Errorf("failed to stat file: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
if info.Mode().Perm() != 0o600 {
|
||||
t.Errorf("file permissions = %o, want %o", info.Mode().Perm(), 0o600)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("http error", func(t *testing.T) {
|
||||
errorServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
w.Write([]byte("not found"))
|
||||
}))
|
||||
defer errorServer.Close()
|
||||
|
||||
localPath := filepath.Join(tmpDir, "error-test", "SKILL.md")
|
||||
err := installer.downloadFile(context.Background(), errorServer.URL, localPath)
|
||||
if err == nil {
|
||||
t.Error("downloadFile() expected error for 404, got nil")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestSkillInstaller_DownloadRaw(t *testing.T) {
|
||||
content := "raw skill content"
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte(content))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
tmpDir := t.TempDir()
|
||||
installer, err := NewSkillInstaller(tmpDir, "", "")
|
||||
if err != nil {
|
||||
t.Fatalf("NewSkillInstaller() error = %v", err)
|
||||
}
|
||||
|
||||
// Replace the client with one that points to our test server
|
||||
// We need to modify the URL in the function, so we'll test indirectly
|
||||
|
||||
localDir := filepath.Join(tmpDir, "raw-test")
|
||||
ctx := context.Background()
|
||||
|
||||
// Create a simple test by calling downloadFile directly since downloadRaw
|
||||
// constructs its own URL
|
||||
testFile := filepath.Join(localDir, "SKILL.md")
|
||||
err = installer.downloadFile(ctx, server.URL, testFile)
|
||||
if err != nil {
|
||||
t.Errorf("downloadFile() error = %v", err)
|
||||
}
|
||||
|
||||
// Verify file content
|
||||
data, err := os.ReadFile(testFile)
|
||||
if err != nil {
|
||||
t.Errorf("failed to read file: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
if string(data) != content {
|
||||
t.Errorf("content = %q, want %q", string(data), content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSkillInstaller_Uninstall(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
skillsDir := filepath.Join(tmpDir, "skills")
|
||||
os.MkdirAll(skillsDir, 0o755)
|
||||
|
||||
installer, err := NewSkillInstaller(tmpDir, "", "")
|
||||
if err != nil {
|
||||
t.Fatalf("NewSkillInstaller() error = %v", err)
|
||||
}
|
||||
|
||||
t.Run("uninstall existing skill", func(t *testing.T) {
|
||||
skillName := "test-skill"
|
||||
skillDir := filepath.Join(skillsDir, skillName)
|
||||
|
||||
// Create skill directory with a file
|
||||
os.MkdirAll(skillDir, 0o755)
|
||||
os.WriteFile(filepath.Join(skillDir, "SKILL.md"), []byte("test"), 0o644)
|
||||
|
||||
if err := installer.Uninstall(skillName); err != nil {
|
||||
t.Errorf("Uninstall() error = %v", err)
|
||||
}
|
||||
|
||||
// Verify directory was removed
|
||||
if _, err := os.Stat(skillDir); !os.IsNotExist(err) {
|
||||
t.Error("skill directory still exists after uninstall")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("uninstall non-existent skill", func(t *testing.T) {
|
||||
if err := installer.Uninstall("non-existent-skill"); err == nil {
|
||||
t.Error("Uninstall() expected error for non-existent skill, got nil")
|
||||
} else if !strings.Contains(err.Error(), "not found") {
|
||||
t.Errorf("error message = %q, want 'not found'", err.Error())
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("uninstall with path separator", func(t *testing.T) {
|
||||
skillName := "owner/repo/skill-name"
|
||||
skillDir := filepath.Join(skillsDir, "skill-name")
|
||||
|
||||
// Create skill directory
|
||||
os.MkdirAll(skillDir, 0o755)
|
||||
os.WriteFile(filepath.Join(skillDir, "SKILL.md"), []byte("test"), 0o644)
|
||||
|
||||
if err := installer.Uninstall(skillName); err != nil {
|
||||
t.Errorf("Uninstall() error = %v", err)
|
||||
}
|
||||
|
||||
if _, err := os.Stat(skillDir); !os.IsNotExist(err) {
|
||||
t.Error("skill directory still exists after uninstall")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("uninstall with trailing slash", func(t *testing.T) {
|
||||
skillName := "skill-name/"
|
||||
skillDir := filepath.Join(skillsDir, "skill-name")
|
||||
|
||||
// Create skill directory
|
||||
os.MkdirAll(skillDir, 0o755)
|
||||
os.WriteFile(filepath.Join(skillDir, "SKILL.md"), []byte("test"), 0o644)
|
||||
|
||||
if err := installer.Uninstall(skillName); err != nil {
|
||||
t.Errorf("Uninstall() error = %v", err)
|
||||
}
|
||||
|
||||
if _, err := os.Stat(skillDir); !os.IsNotExist(err) {
|
||||
t.Error("skill directory still exists after uninstall")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestSkillInstaller_InstallFromGitHub_SkillAlreadyExists(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
skillsDir := filepath.Join(tmpDir, "skills")
|
||||
os.MkdirAll(skillsDir, 0o755)
|
||||
|
||||
installer, err := NewSkillInstaller(tmpDir, "", "")
|
||||
if err != nil {
|
||||
t.Fatalf("NewSkillInstaller() error = %v", err)
|
||||
}
|
||||
|
||||
// Create an existing skill directory
|
||||
existingSkill := filepath.Join(skillsDir, "picoclaw")
|
||||
os.MkdirAll(existingSkill, 0o755)
|
||||
os.WriteFile(filepath.Join(existingSkill, "SKILL.md"), []byte("existing"), 0o644)
|
||||
|
||||
// Try to install the same skill - should fail
|
||||
err = installer.InstallFromGitHub(context.Background(), "sipeed/picoclaw")
|
||||
if err == nil {
|
||||
t.Error("InstallFromGitHub() expected error for existing skill, got nil")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "already exists") {
|
||||
t.Errorf("error message = %q, want 'already exists'", err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
func TestGitHubContent_Struct(t *testing.T) {
|
||||
// Test that GitHubContent struct can be properly unmarshaled
|
||||
jsonData := `{
|
||||
"name": "test.md",
|
||||
"path": "skills/test.md",
|
||||
"type": "file",
|
||||
"download_url": "https://example.com/download",
|
||||
"url": "https://api.github.com/contents/skills/test.md"
|
||||
}`
|
||||
|
||||
var content GitHubContent
|
||||
err := json.Unmarshal([]byte(jsonData), &content)
|
||||
if err != nil {
|
||||
t.Errorf("failed to unmarshal GitHubContent: %v", err)
|
||||
}
|
||||
|
||||
if content.Name != "test.md" {
|
||||
t.Errorf("Name = %q, want 'test.md'", content.Name)
|
||||
}
|
||||
if content.Type != "file" {
|
||||
t.Errorf("Type = %q, want 'file'", content.Type)
|
||||
}
|
||||
if content.DownloadURL != "https://example.com/download" {
|
||||
t.Errorf("DownloadURL = %q, want 'https://example.com/download'", content.DownloadURL)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSkillInstaller_GetGithubDirAllFiles(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
installer, err := NewSkillInstaller(tmpDir, "", "")
|
||||
if err != nil {
|
||||
t.Fatalf("NewSkillInstaller() error = %v", err)
|
||||
}
|
||||
|
||||
// Create a test server that mimics GitHub API
|
||||
fileContent := "skill file content"
|
||||
var serverURL string
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
// Check for authorization header
|
||||
authHeader := r.Header.Get("Authorization")
|
||||
if authHeader != "" && !strings.HasPrefix(authHeader, "Bearer ") {
|
||||
t.Errorf("expected Bearer token, got: %s", authHeader)
|
||||
}
|
||||
|
||||
// Return different responses based on path
|
||||
if strings.Contains(r.URL.Path, "/contents") {
|
||||
// API response for directory listing
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
|
||||
items := []map[string]any{
|
||||
{
|
||||
"name": "SKILL.md",
|
||||
"path": "SKILL.md",
|
||||
"type": "file",
|
||||
"download_url": serverURL + "/download/SKILL.md",
|
||||
},
|
||||
{
|
||||
"name": "scripts",
|
||||
"path": "scripts",
|
||||
"type": "dir",
|
||||
"url": serverURL + "/api/scripts",
|
||||
},
|
||||
}
|
||||
json.NewEncoder(w).Encode(items)
|
||||
} else if strings.Contains(r.URL.Path, "/api/scripts") {
|
||||
// API response for scripts subdirectory
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
|
||||
items := []map[string]any{
|
||||
{
|
||||
"name": "test.py",
|
||||
"path": "scripts/test.py",
|
||||
"type": "file",
|
||||
"download_url": serverURL + "/download/test.py",
|
||||
},
|
||||
}
|
||||
json.NewEncoder(w).Encode(items)
|
||||
} else if strings.Contains(r.URL.Path, "/download/") {
|
||||
// Raw file download
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte(fileContent))
|
||||
} else {
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
}
|
||||
}))
|
||||
serverURL = server.URL
|
||||
defer server.Close()
|
||||
|
||||
localDir := filepath.Join(tmpDir, "test-skill")
|
||||
|
||||
t.Run("download from GitHub API", func(t *testing.T) {
|
||||
err := installer.getGithubDirAllFiles(context.Background(), server.URL+"/contents", localDir, true)
|
||||
if err != nil {
|
||||
t.Errorf("getGithubDirAllFiles() error = %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
// Verify SKILL.md was downloaded
|
||||
skillMd := filepath.Join(localDir, "SKILL.md")
|
||||
data, err := os.ReadFile(skillMd)
|
||||
if err != nil {
|
||||
t.Errorf("failed to read SKILL.md: %v", err)
|
||||
return
|
||||
}
|
||||
if string(data) != fileContent {
|
||||
t.Errorf("SKILL.md content = %q, want %q", string(data), fileContent)
|
||||
}
|
||||
|
||||
// Verify scripts directory and file
|
||||
scriptFile := filepath.Join(localDir, "scripts", "test.py")
|
||||
data, err = os.ReadFile(scriptFile)
|
||||
if err != nil {
|
||||
t.Errorf("failed to read test.py: %v", err)
|
||||
return
|
||||
}
|
||||
if string(data) != fileContent {
|
||||
t.Errorf("test.py content = %q, want %q", string(data), fileContent)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("http error response", func(t *testing.T) {
|
||||
errorServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusForbidden)
|
||||
}))
|
||||
defer errorServer.Close()
|
||||
|
||||
err := installer.getGithubDirAllFiles(
|
||||
context.Background(),
|
||||
errorServer.URL,
|
||||
filepath.Join(tmpDir, "error-test"),
|
||||
true,
|
||||
)
|
||||
if err == nil {
|
||||
t.Error("getGithubDirAllFiles() expected error for 403, got nil")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestSkillInstaller_InstallFromGitHub_WithToken(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
skillsDir := filepath.Join(tmpDir, "skills")
|
||||
os.MkdirAll(skillsDir, 0o755)
|
||||
|
||||
var serverURL string
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
// Capture the authorization header
|
||||
authHeader := r.Header.Get("Authorization")
|
||||
if authHeader != "" {
|
||||
tokenReceived := strings.TrimPrefix(authHeader, "Bearer ")
|
||||
t.Fatalf("github token is %s", tokenReceived)
|
||||
}
|
||||
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
|
||||
items := []map[string]any{
|
||||
{
|
||||
"name": "SKILL.md",
|
||||
"path": "SKILL.md",
|
||||
"type": "file",
|
||||
"download_url": serverURL + "/download/SKILL.md",
|
||||
},
|
||||
}
|
||||
json.NewEncoder(w).Encode(items)
|
||||
}))
|
||||
serverURL = server.URL
|
||||
defer server.Close()
|
||||
|
||||
installer, err := NewSkillInstaller(tmpDir, "test-github-token", "")
|
||||
if err != nil {
|
||||
t.Fatalf("NewSkillInstaller() error = %v", err)
|
||||
}
|
||||
|
||||
// We need to test the token is passed - the actual install will fail
|
||||
// because we're not fully mocking the download, but we can verify
|
||||
// the token is sent in the request
|
||||
|
||||
// Use a simple context with timeout
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
// The install will fail because download URL isn't properly set up,
|
||||
// but the token should be sent in the API request
|
||||
_ = installer.InstallFromGitHub(ctx, "owner/repo")
|
||||
|
||||
// Note: We can't easily intercept the download request since it's a different URL,
|
||||
// but the fact that the API request was made verifies the token flow
|
||||
// In a real scenario, the token would be sent to both API and raw downloads
|
||||
}
|
||||
|
||||
func TestSkillInstaller_ContextCancellation(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
installer, err := NewSkillInstaller(tmpDir, "", "")
|
||||
if err != nil {
|
||||
t.Fatalf("NewSkillInstaller() error = %v", err)
|
||||
}
|
||||
|
||||
// Create a slow server
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte("response"))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
// Create a canceled context
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel() // Cancel immediately
|
||||
|
||||
localPath := filepath.Join(tmpDir, "cancel-test", "file.txt")
|
||||
err = installer.downloadFile(ctx, server.URL, localPath)
|
||||
|
||||
if err == nil {
|
||||
t.Error("downloadFile() expected error for canceled context, got nil")
|
||||
}
|
||||
}
|
||||
|
|
@ -14,6 +14,8 @@ import (
|
|||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/utils"
|
||||
)
|
||||
|
||||
const (
|
||||
|
|
@ -41,43 +43,6 @@ var (
|
|||
reDDGSnippet = regexp.MustCompile(`<a class="result__snippet[^"]*".*?>([\s\S]*?)</a>`)
|
||||
)
|
||||
|
||||
// createHTTPClient creates an HTTP client with optional proxy support
|
||||
func createHTTPClient(proxyURL string, timeout time.Duration) (*http.Client, error) {
|
||||
client := &http.Client{
|
||||
Timeout: timeout,
|
||||
Transport: &http.Transport{
|
||||
MaxIdleConns: 10,
|
||||
IdleConnTimeout: 30 * time.Second,
|
||||
DisableCompression: false,
|
||||
TLSHandshakeTimeout: 15 * time.Second,
|
||||
},
|
||||
}
|
||||
|
||||
if proxyURL != "" {
|
||||
proxy, err := url.Parse(proxyURL)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid proxy URL: %w", err)
|
||||
}
|
||||
scheme := strings.ToLower(proxy.Scheme)
|
||||
switch scheme {
|
||||
case "http", "https", "socks5", "socks5h":
|
||||
default:
|
||||
return nil, fmt.Errorf(
|
||||
"unsupported proxy scheme %q (supported: http, https, socks5, socks5h)",
|
||||
proxy.Scheme,
|
||||
)
|
||||
}
|
||||
if proxy.Host == "" {
|
||||
return nil, fmt.Errorf("invalid proxy URL: missing host")
|
||||
}
|
||||
client.Transport.(*http.Transport).Proxy = http.ProxyURL(proxy)
|
||||
} else {
|
||||
client.Transport.(*http.Transport).Proxy = http.ProxyFromEnvironment
|
||||
}
|
||||
|
||||
return client, nil
|
||||
}
|
||||
|
||||
type APIKeyPool struct {
|
||||
keys []string
|
||||
current uint32
|
||||
|
|
@ -678,7 +643,7 @@ func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) {
|
|||
maxResults := 5
|
||||
// Priority: Perplexity > Brave > SearXNG > Tavily > DuckDuckGo > GLM Search
|
||||
if opts.PerplexityEnabled && len(opts.PerplexityAPIKeys) > 0 {
|
||||
client, err := createHTTPClient(opts.Proxy, perplexityTimeout)
|
||||
client, err := utils.CreateHTTPClient(opts.Proxy, perplexityTimeout)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create HTTP client for Perplexity: %w", err)
|
||||
}
|
||||
|
|
@ -691,7 +656,7 @@ func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) {
|
|||
maxResults = opts.PerplexityMaxResults
|
||||
}
|
||||
} else if opts.BraveEnabled && len(opts.BraveAPIKeys) > 0 {
|
||||
client, err := createHTTPClient(opts.Proxy, searchTimeout)
|
||||
client, err := utils.CreateHTTPClient(opts.Proxy, searchTimeout)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create HTTP client for Brave: %w", err)
|
||||
}
|
||||
|
|
@ -705,7 +670,7 @@ func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) {
|
|||
maxResults = opts.SearXNGMaxResults
|
||||
}
|
||||
} else if opts.TavilyEnabled && len(opts.TavilyAPIKeys) > 0 {
|
||||
client, err := createHTTPClient(opts.Proxy, searchTimeout)
|
||||
client, err := utils.CreateHTTPClient(opts.Proxy, searchTimeout)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create HTTP client for Tavily: %w", err)
|
||||
}
|
||||
|
|
@ -719,7 +684,7 @@ func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) {
|
|||
maxResults = opts.TavilyMaxResults
|
||||
}
|
||||
} else if opts.DuckDuckGoEnabled {
|
||||
client, err := createHTTPClient(opts.Proxy, searchTimeout)
|
||||
client, err := utils.CreateHTTPClient(opts.Proxy, searchTimeout)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create HTTP client for DuckDuckGo: %w", err)
|
||||
}
|
||||
|
|
@ -728,7 +693,7 @@ func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) {
|
|||
maxResults = opts.DuckDuckGoMaxResults
|
||||
}
|
||||
} else if opts.GLMSearchEnabled && opts.GLMSearchAPIKey != "" {
|
||||
client, err := createHTTPClient(opts.Proxy, searchTimeout)
|
||||
client, err := utils.CreateHTTPClient(opts.Proxy, searchTimeout)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create HTTP client for GLM Search: %w", err)
|
||||
}
|
||||
|
|
@ -827,7 +792,7 @@ func NewWebFetchToolWithProxy(maxChars int, proxy string, fetchLimitBytes int64)
|
|||
if maxChars <= 0 {
|
||||
maxChars = defaultMaxChars
|
||||
}
|
||||
client, err := createHTTPClient(proxy, fetchTimeout)
|
||||
client, err := utils.CreateHTTPClient(proxy, fetchTimeout)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create HTTP client for web fetch: %w", err)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -10,7 +10,6 @@ import (
|
|||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/logger"
|
||||
)
|
||||
|
|
@ -639,108 +638,6 @@ func TestWebTool_WebFetch_MissingDomain(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestCreateHTTPClient_ProxyConfigured(t *testing.T) {
|
||||
client, err := createHTTPClient("http://127.0.0.1:7890", 12*time.Second)
|
||||
if err != nil {
|
||||
t.Fatalf("createHTTPClient() error: %v", err)
|
||||
}
|
||||
if client.Timeout != 12*time.Second {
|
||||
t.Fatalf("client.Timeout = %v, want %v", client.Timeout, 12*time.Second)
|
||||
}
|
||||
|
||||
tr, ok := client.Transport.(*http.Transport)
|
||||
if !ok {
|
||||
t.Fatalf("client.Transport type = %T, want *http.Transport", client.Transport)
|
||||
}
|
||||
if tr.Proxy == nil {
|
||||
t.Fatal("transport.Proxy is nil, want non-nil")
|
||||
}
|
||||
|
||||
req, err := http.NewRequest("GET", "https://example.com", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("http.NewRequest() error: %v", err)
|
||||
}
|
||||
proxyURL, err := tr.Proxy(req)
|
||||
if err != nil {
|
||||
t.Fatalf("transport.Proxy(req) error: %v", err)
|
||||
}
|
||||
if proxyURL == nil || proxyURL.String() != "http://127.0.0.1:7890" {
|
||||
t.Fatalf("proxy URL = %v, want %q", proxyURL, "http://127.0.0.1:7890")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateHTTPClient_InvalidProxy(t *testing.T) {
|
||||
_, err := createHTTPClient("://bad-proxy", 10*time.Second)
|
||||
if err == nil {
|
||||
t.Fatal("createHTTPClient() expected error for invalid proxy URL, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateHTTPClient_Socks5ProxyConfigured(t *testing.T) {
|
||||
client, err := createHTTPClient("socks5://127.0.0.1:1080", 8*time.Second)
|
||||
if err != nil {
|
||||
t.Fatalf("createHTTPClient() error: %v", err)
|
||||
}
|
||||
|
||||
tr, ok := client.Transport.(*http.Transport)
|
||||
if !ok {
|
||||
t.Fatalf("client.Transport type = %T, want *http.Transport", client.Transport)
|
||||
}
|
||||
req, err := http.NewRequest("GET", "https://example.com", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("http.NewRequest() error: %v", err)
|
||||
}
|
||||
proxyURL, err := tr.Proxy(req)
|
||||
if err != nil {
|
||||
t.Fatalf("transport.Proxy(req) error: %v", err)
|
||||
}
|
||||
if proxyURL == nil || proxyURL.String() != "socks5://127.0.0.1:1080" {
|
||||
t.Fatalf("proxy URL = %v, want %q", proxyURL, "socks5://127.0.0.1:1080")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateHTTPClient_UnsupportedProxyScheme(t *testing.T) {
|
||||
_, err := createHTTPClient("ftp://127.0.0.1:21", 10*time.Second)
|
||||
if err == nil {
|
||||
t.Fatal("createHTTPClient() expected error for unsupported scheme, got nil")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "unsupported proxy scheme") {
|
||||
t.Fatalf("error = %q, want to contain %q", err.Error(), "unsupported proxy scheme")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateHTTPClient_ProxyFromEnvironmentWhenConfigEmpty(t *testing.T) {
|
||||
t.Setenv("HTTP_PROXY", "http://127.0.0.1:8888")
|
||||
t.Setenv("http_proxy", "http://127.0.0.1:8888")
|
||||
t.Setenv("HTTPS_PROXY", "http://127.0.0.1:8888")
|
||||
t.Setenv("https_proxy", "http://127.0.0.1:8888")
|
||||
t.Setenv("ALL_PROXY", "")
|
||||
t.Setenv("all_proxy", "")
|
||||
t.Setenv("NO_PROXY", "")
|
||||
t.Setenv("no_proxy", "")
|
||||
|
||||
client, err := createHTTPClient("", 10*time.Second)
|
||||
if err != nil {
|
||||
t.Fatalf("createHTTPClient() error: %v", err)
|
||||
}
|
||||
|
||||
tr, ok := client.Transport.(*http.Transport)
|
||||
if !ok {
|
||||
t.Fatalf("client.Transport type = %T, want *http.Transport", client.Transport)
|
||||
}
|
||||
if tr.Proxy == nil {
|
||||
t.Fatal("transport.Proxy is nil, want proxy function from environment")
|
||||
}
|
||||
|
||||
req, err := http.NewRequest("GET", "https://example.com", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("http.NewRequest() error: %v", err)
|
||||
}
|
||||
if _, err := tr.Proxy(req); err != nil {
|
||||
t.Fatalf("transport.Proxy(req) error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewWebFetchToolWithProxy(t *testing.T) {
|
||||
tool, err := NewWebFetchToolWithProxy(1024, "http://127.0.0.1:7890", testFetchLimit)
|
||||
if err != nil {
|
||||
|
|
|
|||
48
pkg/utils/http_client.go
Normal file
48
pkg/utils/http_client.go
Normal file
|
|
@ -0,0 +1,48 @@
|
|||
package utils
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// CreateHTTPClient creates an HTTP client with optional proxy support.
|
||||
// If proxyURL is empty, it uses the system environment proxy settings.
|
||||
// Supported proxy schemes: http, https, socks5, socks5h.
|
||||
func CreateHTTPClient(proxyURL string, timeout time.Duration) (*http.Client, error) {
|
||||
client := &http.Client{
|
||||
Timeout: timeout,
|
||||
Transport: &http.Transport{
|
||||
MaxIdleConns: 10,
|
||||
IdleConnTimeout: 30 * time.Second,
|
||||
DisableCompression: false,
|
||||
TLSHandshakeTimeout: 15 * time.Second,
|
||||
},
|
||||
}
|
||||
|
||||
if proxyURL != "" {
|
||||
proxy, err := url.Parse(proxyURL)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid proxy URL: %w", err)
|
||||
}
|
||||
scheme := strings.ToLower(proxy.Scheme)
|
||||
switch scheme {
|
||||
case "http", "https", "socks5", "socks5h":
|
||||
default:
|
||||
return nil, fmt.Errorf(
|
||||
"unsupported proxy scheme %q (supported: http, https, socks5, socks5h)",
|
||||
proxy.Scheme,
|
||||
)
|
||||
}
|
||||
if proxy.Host == "" {
|
||||
return nil, fmt.Errorf("invalid proxy URL: missing host")
|
||||
}
|
||||
client.Transport.(*http.Transport).Proxy = http.ProxyURL(proxy)
|
||||
} else {
|
||||
client.Transport.(*http.Transport).Proxy = http.ProxyFromEnvironment
|
||||
}
|
||||
|
||||
return client, nil
|
||||
}
|
||||
110
pkg/utils/http_client_test.go
Normal file
110
pkg/utils/http_client_test.go
Normal file
|
|
@ -0,0 +1,110 @@
|
|||
package utils
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestCreateHTTPClient_ProxyConfigured(t *testing.T) {
|
||||
client, err := CreateHTTPClient("http://127.0.0.1:7890", 12*time.Second)
|
||||
if err != nil {
|
||||
t.Fatalf("createHTTPClient() error: %v", err)
|
||||
}
|
||||
if client.Timeout != 12*time.Second {
|
||||
t.Fatalf("client.Timeout = %v, want %v", client.Timeout, 12*time.Second)
|
||||
}
|
||||
|
||||
tr, ok := client.Transport.(*http.Transport)
|
||||
if !ok {
|
||||
t.Fatalf("client.Transport type = %T, want *http.Transport", client.Transport)
|
||||
}
|
||||
if tr.Proxy == nil {
|
||||
t.Fatal("transport.Proxy is nil, want non-nil")
|
||||
}
|
||||
|
||||
req, err := http.NewRequest("GET", "https://example.com", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("http.NewRequest() error: %v", err)
|
||||
}
|
||||
proxyURL, err := tr.Proxy(req)
|
||||
if err != nil {
|
||||
t.Fatalf("transport.Proxy(req) error: %v", err)
|
||||
}
|
||||
if proxyURL == nil || proxyURL.String() != "http://127.0.0.1:7890" {
|
||||
t.Fatalf("proxy URL = %v, want %q", proxyURL, "http://127.0.0.1:7890")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateHTTPClient_InvalidProxy(t *testing.T) {
|
||||
_, err := CreateHTTPClient("://bad-proxy", 10*time.Second)
|
||||
if err == nil {
|
||||
t.Fatal("createHTTPClient() expected error for invalid proxy URL, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateHTTPClient_Socks5ProxyConfigured(t *testing.T) {
|
||||
client, err := CreateHTTPClient("socks5://127.0.0.1:1080", 8*time.Second)
|
||||
if err != nil {
|
||||
t.Fatalf("createHTTPClient() error: %v", err)
|
||||
}
|
||||
|
||||
tr, ok := client.Transport.(*http.Transport)
|
||||
if !ok {
|
||||
t.Fatalf("client.Transport type = %T, want *http.Transport", client.Transport)
|
||||
}
|
||||
req, err := http.NewRequest("GET", "https://example.com", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("http.NewRequest() error: %v", err)
|
||||
}
|
||||
proxyURL, err := tr.Proxy(req)
|
||||
if err != nil {
|
||||
t.Fatalf("transport.Proxy(req) error: %v", err)
|
||||
}
|
||||
if proxyURL == nil || proxyURL.String() != "socks5://127.0.0.1:1080" {
|
||||
t.Fatalf("proxy URL = %v, want %q", proxyURL, "socks5://127.0.0.1:1080")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateHTTPClient_UnsupportedProxyScheme(t *testing.T) {
|
||||
_, err := CreateHTTPClient("ftp://127.0.0.1:21", 10*time.Second)
|
||||
if err == nil {
|
||||
t.Fatal("createHTTPClient() expected error for unsupported scheme, got nil")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "unsupported proxy scheme") {
|
||||
t.Fatalf("error = %q, want to contain %q", err.Error(), "unsupported proxy scheme")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateHTTPClient_ProxyFromEnvironmentWhenConfigEmpty(t *testing.T) {
|
||||
t.Setenv("HTTP_PROXY", "http://127.0.0.1:8888")
|
||||
t.Setenv("http_proxy", "http://127.0.0.1:8888")
|
||||
t.Setenv("HTTPS_PROXY", "http://127.0.0.1:8888")
|
||||
t.Setenv("https_proxy", "http://127.0.0.1:8888")
|
||||
t.Setenv("ALL_PROXY", "")
|
||||
t.Setenv("all_proxy", "")
|
||||
t.Setenv("NO_PROXY", "")
|
||||
t.Setenv("no_proxy", "")
|
||||
|
||||
client, err := CreateHTTPClient("", 10*time.Second)
|
||||
if err != nil {
|
||||
t.Fatalf("createHTTPClient() error: %v", err)
|
||||
}
|
||||
|
||||
tr, ok := client.Transport.(*http.Transport)
|
||||
if !ok {
|
||||
t.Fatalf("client.Transport type = %T, want *http.Transport", client.Transport)
|
||||
}
|
||||
if tr.Proxy == nil {
|
||||
t.Fatal("transport.Proxy is nil, want proxy function from environment")
|
||||
}
|
||||
|
||||
req, err := http.NewRequest("GET", "https://example.com", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("http.NewRequest() error: %v", err)
|
||||
}
|
||||
if _, err := tr.Proxy(req); err != nil {
|
||||
t.Fatalf("transport.Proxy(req) error: %v", err)
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue