diff --git a/.github/workflows/pr.yml b/.github/workflows/pr.yml index 902d4d4eb..2d544d4f0 100644 --- a/.github/workflows/pr.yml +++ b/.github/workflows/pr.yml @@ -23,10 +23,13 @@ jobs: uses: golangci/golangci-lint-action@v9 with: version: v2.10.1 + args: --build-tags=goolm,stdjson vuln_check: name: Security Check runs-on: ubuntu-latest + env: + GOFLAGS: -tags=goolm,stdjson steps: - name: Checkout uses: actions/checkout@v6 @@ -59,4 +62,4 @@ jobs: run: go generate ./... - name: Run go test - run: go test ./... + run: go test -tags goolm,stdjson ./... diff --git a/.gitignore b/.gitignore index 72f3b1761..b869ecc33 100644 --- a/.gitignore +++ b/.gitignore @@ -25,6 +25,9 @@ build/ # Secrets & Config (keep templates, ignore actual secrets) .env config/config.json +.security.yml +onboard + # Test coverage.txt diff --git a/.goreleaser.yaml b/.goreleaser.yaml index a73f87f30..9c26de34f 100644 --- a/.goreleaser.yaml +++ b/.goreleaser.yaml @@ -2,6 +2,11 @@ # vim: set ts=2 sw=2 tw=0 fo=cnqoj version: 2 +git: + ignore_tags: + - nightly + - ".*-nightly.*" + before: hooks: - go mod tidy @@ -15,6 +20,7 @@ builds: env: - CGO_ENABLED=0 tags: + - goolm - stdjson ldflags: - -s -w @@ -57,6 +63,7 @@ builds: env: - CGO_ENABLED=0 tags: + - goolm - stdjson ldflags: - -s -w @@ -95,6 +102,7 @@ builds: env: - CGO_ENABLED=0 tags: + - goolm - stdjson ldflags: - -s -w diff --git a/Makefile b/Makefile index 411cd9dc5..9581fa633 100644 --- a/Makefile +++ b/Makefile @@ -17,7 +17,13 @@ LDFLAGS=-X $(CONFIG_PKG).Version=$(VERSION) -X $(CONFIG_PKG).GitCommit=$(GIT_COM # Go variables GO?=CGO_ENABLED=0 go WEB_GO?=$(GO) -GOFLAGS?=-v -tags stdjson +GO_BUILD_TAGS?=goolm,stdjson +GOFLAGS?=-v -tags $(GO_BUILD_TAGS) +comma:=, +empty:= +space:=$(empty) $(empty) +GO_BUILD_TAGS_NO_GOOLM:=$(subst $(space),$(comma),$(strip $(filter-out goolm,$(subst $(comma),$(space),$(GO_BUILD_TAGS))))) +GOFLAGS_NO_GOOLM?=-v -tags $(GO_BUILD_TAGS_NO_GOOLM) # Patch MIPS LE ELF e_flags (offset 36) for NaN2008-only kernels (e.g. Ingenic X2600). # @@ -41,6 +47,13 @@ define PATCH_MIPS_FLAGS fi endef +# Patch creack/pty for loong64 support (upstream doesn't have ztypes_loong64.go) +PTY_PATCH_LOONG64=pty_dir=$$(go env GOMODCACHE)/github.com/creack/pty@v1.1.9; \ + if [ -d "$$pty_dir" ] && [ ! -f "$$pty_dir/ztypes_loong64.go" ]; then \ + chmod +w "$$pty_dir" 2>/dev/null || true; \ + printf '//go:build linux && loong64\npackage pty\ntype (_C_int int32; _C_uint uint32)\n' > "$$pty_dir/ztypes_loong64.go"; \ + fi + # Golangci-lint GOLANGCI_LINT?=golangci-lint @@ -125,20 +138,28 @@ build-launcher: @ln -sf picoclaw-launcher-$(PLATFORM)-$(ARCH) $(BUILD_DIR)/picoclaw-launcher @echo "Build complete: $(BUILD_DIR)/picoclaw-launcher" +## build-launcher-tui: Build the picoclaw-launcher TUI binary +build-launcher-tui: + @echo "Building picoclaw-launcher-tui for $(PLATFORM)/$(ARCH)..." + @mkdir -p $(BUILD_DIR) + @$(GO) build $(GOFLAGS) -o $(BUILD_DIR)/picoclaw-launcher-tui-$(PLATFORM)-$(ARCH) ./cmd/picoclaw-launcher-tui + @ln -sf picoclaw-launcher-tui-$(PLATFORM)-$(ARCH) $(BUILD_DIR)/picoclaw-launcher-tui + @echo "Build complete: $(BUILD_DIR)/picoclaw-launcher-tui" + ## build-whatsapp-native: Build with WhatsApp native (whatsmeow) support; larger binary build-whatsapp-native: generate ## @echo "Building $(BINARY_NAME) with WhatsApp native for $(PLATFORM)/$(ARCH)..." @echo "Building for multiple platforms..." @mkdir -p $(BUILD_DIR) - GOOS=linux GOARCH=amd64 $(GO) build -tags whatsapp_native -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-amd64 ./$(CMD_DIR) - GOOS=linux GOARCH=arm GOARM=7 $(GO) build -tags whatsapp_native -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-arm ./$(CMD_DIR) - GOOS=linux GOARCH=arm64 $(GO) build -tags whatsapp_native -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-arm64 ./$(CMD_DIR) - GOOS=linux GOARCH=loong64 $(GO) build -tags whatsapp_native -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-loong64 ./$(CMD_DIR) - GOOS=linux GOARCH=riscv64 $(GO) build -tags whatsapp_native -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-riscv64 ./$(CMD_DIR) - GOOS=linux GOARCH=mipsle GOMIPS=softfloat $(GO) build -tags whatsapp_native -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-mipsle ./$(CMD_DIR) + GOOS=linux GOARCH=amd64 $(GO) build -tags $(GO_BUILD_TAGS),whatsapp_native -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-amd64 ./$(CMD_DIR) + GOOS=linux GOARCH=arm GOARM=7 $(GO) build -tags $(GO_BUILD_TAGS),whatsapp_native -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-arm ./$(CMD_DIR) + GOOS=linux GOARCH=arm64 $(GO) build -tags $(GO_BUILD_TAGS),whatsapp_native -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-arm64 ./$(CMD_DIR) + GOOS=linux GOARCH=loong64 $(GO) build -tags $(GO_BUILD_TAGS),whatsapp_native -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-loong64 ./$(CMD_DIR) + GOOS=linux GOARCH=riscv64 $(GO) build -tags $(GO_BUILD_TAGS),whatsapp_native -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-riscv64 ./$(CMD_DIR) + GOOS=linux GOARCH=mipsle GOMIPS=softfloat $(GO) build -tags $(GO_BUILD_TAGS_NO_GOOLM),whatsapp_native -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-mipsle ./$(CMD_DIR) $(call PATCH_MIPS_FLAGS,$(BUILD_DIR)/$(BINARY_NAME)-linux-mipsle) - GOOS=darwin GOARCH=arm64 $(GO) build -tags whatsapp_native -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-darwin-arm64 ./$(CMD_DIR) - GOOS=windows GOARCH=amd64 $(GO) build -tags whatsapp_native -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-windows-amd64.exe ./$(CMD_DIR) + GOOS=darwin GOARCH=arm64 $(GO) build -tags $(GO_BUILD_TAGS),whatsapp_native -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-darwin-arm64 ./$(CMD_DIR) + GOOS=windows GOARCH=amd64 $(GO) build -tags $(GO_BUILD_TAGS),whatsapp_native -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-windows-amd64.exe ./$(CMD_DIR) ## @$(GO) build $(GOFLAGS) -tags whatsapp_native -ldflags "$(LDFLAGS)" -o $(BINARY_PATH) ./$(CMD_DIR) @echo "Build complete" ## @ln -sf $(BINARY_NAME)-$(PLATFORM)-$(ARCH) $(BUILD_DIR)/$(BINARY_NAME) @@ -147,21 +168,21 @@ build-whatsapp-native: generate build-linux-arm: generate @echo "Building for linux/arm (GOARM=7)..." @mkdir -p $(BUILD_DIR) - GOOS=linux GOARCH=arm GOARM=7 $(GO) build -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-arm ./$(CMD_DIR) + GOOS=linux GOARCH=arm GOARM=7 $(GO) build $(GOFLAGS) -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-arm ./$(CMD_DIR) @echo "Build complete: $(BUILD_DIR)/$(BINARY_NAME)-linux-arm" ## build-linux-arm64: Build for Linux ARM64 (e.g. Raspberry Pi Zero 2 W 64-bit) build-linux-arm64: generate @echo "Building for linux/arm64..." @mkdir -p $(BUILD_DIR) - GOOS=linux GOARCH=arm64 $(GO) build -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-arm64 ./$(CMD_DIR) + GOOS=linux GOARCH=arm64 $(GO) build $(GOFLAGS) -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-arm64 ./$(CMD_DIR) @echo "Build complete: $(BUILD_DIR)/$(BINARY_NAME)-linux-arm64" ## build-linux-mipsle: Build for Linux MIPS32 LE build-linux-mipsle: generate @echo "Building for linux/mipsle (softfloat)..." @mkdir -p $(BUILD_DIR) - GOOS=linux GOARCH=mipsle GOMIPS=softfloat $(GO) build -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-mipsle ./$(CMD_DIR) + GOOS=linux GOARCH=mipsle GOMIPS=softfloat $(GO) build $(GOFLAGS_NO_GOOLM) -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-mipsle ./$(CMD_DIR) $(call PATCH_MIPS_FLAGS,$(BUILD_DIR)/$(BINARY_NAME)-linux-mipsle) @echo "Build complete: $(BUILD_DIR)/$(BINARY_NAME)-linux-mipsle" @@ -173,18 +194,19 @@ build-pi-zero: build-linux-arm build-linux-arm64 build-all: generate @echo "Building for multiple platforms..." @mkdir -p $(BUILD_DIR) - GOOS=linux GOARCH=amd64 $(GO) build -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-amd64 ./$(CMD_DIR) - GOOS=linux GOARCH=arm GOARM=7 $(GO) build -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-arm ./$(CMD_DIR) - GOOS=linux GOARCH=arm64 $(GO) build -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-arm64 ./$(CMD_DIR) - GOOS=linux GOARCH=loong64 $(GO) build -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-loong64 ./$(CMD_DIR) - GOOS=linux GOARCH=riscv64 $(GO) build -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-riscv64 ./$(CMD_DIR) - GOOS=linux GOARCH=mipsle GOMIPS=softfloat $(GO) build -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-mipsle ./$(CMD_DIR) + GOOS=linux GOARCH=amd64 $(GO) build $(GOFLAGS) -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-amd64 ./$(CMD_DIR) + GOOS=linux GOARCH=arm GOARM=7 $(GO) build $(GOFLAGS) -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-arm ./$(CMD_DIR) + GOOS=linux GOARCH=arm64 $(GO) build $(GOFLAGS) -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-arm64 ./$(CMD_DIR) + @$(PTY_PATCH_LOONG64) + GOOS=linux GOARCH=loong64 $(GO) build $(GOFLAGS) -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-loong64 ./$(CMD_DIR) + GOOS=linux GOARCH=riscv64 $(GO) build $(GOFLAGS) -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-riscv64 ./$(CMD_DIR) + GOOS=linux GOARCH=mipsle GOMIPS=softfloat $(GO) build $(GOFLAGS_NO_GOOLM) -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-mipsle ./$(CMD_DIR) $(call PATCH_MIPS_FLAGS,$(BUILD_DIR)/$(BINARY_NAME)-linux-mipsle) - GOOS=linux GOARCH=arm GOARM=7 $(GO) build -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-armv7 ./$(CMD_DIR) - GOOS=darwin GOARCH=arm64 $(GO) build -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-darwin-arm64 ./$(CMD_DIR) - GOOS=windows GOARCH=amd64 $(GO) build -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-windows-amd64.exe ./$(CMD_DIR) - GOOS=netbsd GOARCH=amd64 $(GO) build -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-netbsd-amd64 ./$(CMD_DIR) - GOOS=netbsd GOARCH=arm64 $(GO) build -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-netbsd-arm64 ./$(CMD_DIR) + GOOS=linux GOARCH=arm GOARM=7 $(GO) build $(GOFLAGS) -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-linux-armv7 ./$(CMD_DIR) + GOOS=darwin GOARCH=arm64 $(GO) build $(GOFLAGS) -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-darwin-arm64 ./$(CMD_DIR) + GOOS=windows GOARCH=amd64 $(GO) build $(GOFLAGS) -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-windows-amd64.exe ./$(CMD_DIR) + GOOS=netbsd GOARCH=amd64 $(GO) build $(GOFLAGS) -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-netbsd-amd64 ./$(CMD_DIR) + GOOS=netbsd GOARCH=arm64 $(GO) build $(GOFLAGS) -ldflags "$(LDFLAGS)" -o $(BUILD_DIR)/$(BINARY_NAME)-netbsd-arm64 ./$(CMD_DIR) @echo "All builds complete" ## install: Install picoclaw to system and copy builtin skills @@ -221,13 +243,13 @@ clean: ## vet: Run go vet for static analysis vet: generate - @packages="$$(go list ./...)" && \ - $(GO) vet $$(printf '%s\n' "$$packages" | grep -v '^github.com/sipeed/picoclaw/web/') + @packages="$$($(GO) list $(GOFLAGS) ./...)" && \ + $(GO) vet $(GOFLAGS) $$(printf '%s\n' "$$packages" | grep -v '^github.com/sipeed/picoclaw/web/') @cd web/backend && $(WEB_GO) vet ./... ## test: Test Go code test: generate - @$(GO) test $$(go list ./... | grep -v github.com/sipeed/picoclaw/web/) + @$(GO) test $(GOFLAGS) $$($(GO) list $(GOFLAGS) ./... | grep -v github.com/sipeed/picoclaw/web/) @cd web && make test ## fmt: Format Go code @@ -236,11 +258,11 @@ fmt: ## lint: Run linters lint: - @$(GOLANGCI_LINT) run + @$(GOLANGCI_LINT) run --build-tags $(GO_BUILD_TAGS) ## fix: Fix linting issues fix: - @$(GOLANGCI_LINT) run --fix + @$(GOLANGCI_LINT) run --fix --build-tags $(GO_BUILD_TAGS) ## deps: Download dependencies deps: diff --git a/README.md b/README.md index 568c87e59..f627e261e 100644 --- a/README.md +++ b/README.md @@ -322,14 +322,17 @@ This creates `~/.picoclaw/config.json` and the workspace directory. "model_list": [ { "model_name": "gpt-5.4", - "model": "openai/gpt-5.4", - "api_key": "sk-your-api-key" + "model": "openai/gpt-5.4" + // api_key is now loaded from .security.yml } ] } ``` > See `config/config.example.json` in the repo for a complete configuration template with all available options. +> +> Please note: config.example.json format is version 0, with sensitive codes in it, and will be auto migrated to version 1+, then, the config.json will only store insensitive data, the sensitive codes will be stored in .security.yml, if you need manually modify the codes, please see `docs/security_configuration.md` for more details. + **3. Chat** diff --git a/assets/wechat.png b/assets/wechat.png index effb4dab9..ecce856af 100644 Binary files a/assets/wechat.png and b/assets/wechat.png differ diff --git a/cmd/picoclaw-launcher-tui/README.md b/cmd/picoclaw-launcher-tui/README.md new file mode 100644 index 000000000..a942045a5 --- /dev/null +++ b/cmd/picoclaw-launcher-tui/README.md @@ -0,0 +1,69 @@ +# Picoclaw Launcher TUI + +This directory contains the terminal-based TUI launcher for `picoclaw`. +It provides a lightweight, terminal-native user interface for managing, configuring, and interacting with the core `picoclaw` engine, without requiring a web browser or graphical environment. + +## Architecture + +The TUI launcher is implemented purely in Go with no external runtime dependencies: +* **`main.go`**: Application entry point, handles initialization and main event loop +* **`ui/`**: TUI interface components built on tview + tcell framework: + - `home.go`: Main dashboard with navigation menu + - `schemes.go`: AI model scheme management + - `users.go`: User and API key management for model providers + - `channels.go`: Communication channel (Telegram/Discord/WeChat etc.) configuration editor + - `gateway.go`: PicoClaw gateway daemon lifecycle management (start/stop/status) + - `app.go`: Core TUI application framework and navigation logic + - `models.go`: Data structures and state management +* **`config/`**: Configuration management layer, integrates with the core picoclaw configuration system + +## Getting Started + +### Prerequisites + +* Go 1.25+ +* Terminal with 256-color support (most modern terminals are compatible) + +### Development + +Run the TUI launcher directly in development mode: + +```bash +# From project root +go run ./cmd/picoclaw-launcher-tui + +# Or from this directory +go run . +``` + +### Build + +Build the standalone TUI launcher binary: + +```bash +# From project root (recommended) +make build-launcher-tui + +# Output will be at: +# build/picoclaw-launcher-tui-- +# with symlink build/picoclaw-launcher-tui + +# Or build directly from this directory +go build -o picoclaw-launcher-tui . +``` + +### Key Features + +* 🖥️ Terminal-native interface - works over SSH, on headless servers, and in low-resource environments +* ⚙️ AI model scheme and API key management +* 📱 Communication channel configuration editor (Telegram/Discord/WeChat etc.) +* 🔄 PicoClaw gateway daemon management (start/stop/status monitoring) +* 💬 One-click launch of interactive AI chat session +* 🎯 Keyboard-first design with intuitive shortcuts + +### Other Commands + +```bash +# Run with custom config file path +go run . /path/to/custom/config.json +``` diff --git a/cmd/picoclaw/internal/agent/helpers.go b/cmd/picoclaw/internal/agent/helpers.go index 0af743bb5..23227d56a 100644 --- a/cmd/picoclaw/internal/agent/helpers.go +++ b/cmd/picoclaw/internal/agent/helpers.go @@ -28,6 +28,8 @@ func agentCmd(message, sessionKey, model string, debug bool) error { return fmt.Errorf("error loading config: %w", err) } + logger.ConfigureFromEnv() + if debug { logger.SetLevel(logger.DEBUG) fmt.Println("🔍 Debug mode enabled") diff --git a/cmd/picoclaw/internal/auth/command.go b/cmd/picoclaw/internal/auth/command.go index 149095699..9de083d8d 100644 --- a/cmd/picoclaw/internal/auth/command.go +++ b/cmd/picoclaw/internal/auth/command.go @@ -17,6 +17,7 @@ func NewAuthCommand() *cobra.Command { newStatusCommand(), newModelsCommand(), newWeixinCommand(), + newWeComCommand(), ) return cmd diff --git a/cmd/picoclaw/internal/auth/command_test.go b/cmd/picoclaw/internal/auth/command_test.go index 12f2bc186..3c7f2d3d6 100644 --- a/cmd/picoclaw/internal/auth/command_test.go +++ b/cmd/picoclaw/internal/auth/command_test.go @@ -33,6 +33,7 @@ func TestNewAuthCommand(t *testing.T) { "status", "models", "weixin", + "wecom", } subcommands := cmd.Commands() diff --git a/cmd/picoclaw/internal/auth/wecom.go b/cmd/picoclaw/internal/auth/wecom.go new file mode 100644 index 000000000..8261f5f80 --- /dev/null +++ b/cmd/picoclaw/internal/auth/wecom.go @@ -0,0 +1,407 @@ +package auth + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "os" + "runtime" + "strconv" + "strings" + "time" + + "github.com/mdp/qrterminal/v3" + "github.com/spf13/cobra" + + "github.com/sipeed/picoclaw/cmd/picoclaw/internal" + "github.com/sipeed/picoclaw/pkg/config" +) + +const ( + wecomQRSourceID = "picoclaw" + wecomQRGenerateEndpoint = "https://work.weixin.qq.com/ai/qc/generate" + wecomQRQueryEndpoint = "https://work.weixin.qq.com/ai/qc/query_result" + wecomQRPageEndpoint = "https://work.weixin.qq.com/ai/qc/gen" + wecomQRHTTPTimeout = 15 * time.Second + wecomQRPollInterval = 3 * time.Second + wecomQRPollTimeout = 5 * time.Minute + wecomDefaultWebSocketURL = "wss://openws.work.weixin.qq.com" +) + +type wecomQRScanner func(context.Context, wecomQRFlowOptions) (wecomQRBotInfo, error) + +type wecomQRFlowOptions struct { + HTTPClient *http.Client + GenerateURL string + QueryURL string + QRCodePageURL string + SourceID string + PollInterval time.Duration + PollTimeout time.Duration + Writer io.Writer +} + +type wecomQRBotInfo struct { + BotID string + Secret string +} + +type wecomQRSession struct { + SCode string + AuthURL string +} + +type wecomQRGenerateResponse struct { + ErrCode int `json:"errcode,omitempty"` + ErrMsg string `json:"errmsg,omitempty"` + Data struct { + SCode string `json:"scode"` + AuthURL string `json:"auth_url"` + } `json:"data"` +} + +type wecomQRQueryResponse struct { + ErrCode int `json:"errcode,omitempty"` + ErrMsg string `json:"errmsg,omitempty"` + Data struct { + Status string `json:"status"` + BotInfo struct { + BotID string `json:"botid"` + Secret string `json:"secret"` + } `json:"bot_info"` + } `json:"data"` +} + +func newWeComCommand() *cobra.Command { + var timeout time.Duration + + cmd := &cobra.Command{ + Use: "wecom", + Short: "Scan a WeCom QR code and configure channels.wecom", + Args: cobra.NoArgs, + RunE: func(_ *cobra.Command, _ []string) error { + return authWeComCmd(timeout) + }, + } + + cmd.Flags().DurationVar(&timeout, "timeout", wecomQRPollTimeout, "How long to wait for QR confirmation") + + return cmd +} + +func authWeComCmd(timeout time.Duration) error { + return authWeComCmdWithScanner(context.Background(), os.Stdout, timeout, scanWeComQRCodeInteractive) +} + +func authWeComCmdWithScanner( + ctx context.Context, + writer io.Writer, + timeout time.Duration, + scanner wecomQRScanner, +) error { + if scanner == nil { + return fmt.Errorf("wecom QR scanner is nil") + } + if writer == nil { + writer = os.Stdout + } + + cfg, err := internal.LoadConfig() + if err != nil { + return fmt.Errorf("failed to load config: %w", err) + } + + opts := defaultWeComQRFlowOptions(timeout) + opts.Writer = writer + + botInfo, err := scanner(ctx, opts) + if err != nil { + return err + } + + applyWeComAuthResult(cfg, botInfo) + + if saveErr := config.SaveConfig(internal.GetConfigPath(), cfg); saveErr != nil { + return fmt.Errorf("failed to save config: %w", saveErr) + } + + fmt.Fprintln(writer) + fmt.Fprintln(writer, "WeCom connected.") + fmt.Fprintf(writer, "Bot ID: %s\n", botInfo.BotID) + fmt.Fprintf(writer, "Config: %s\n", internal.GetConfigPath()) + + return nil +} + +func defaultWeComQRFlowOptions(timeout time.Duration) wecomQRFlowOptions { + if timeout <= 0 { + timeout = wecomQRPollTimeout + } + + return wecomQRFlowOptions{ + HTTPClient: &http.Client{Timeout: wecomQRHTTPTimeout}, + GenerateURL: wecomQRGenerateEndpoint, + QueryURL: wecomQRQueryEndpoint, + QRCodePageURL: wecomQRPageEndpoint, + SourceID: wecomQRSourceID, + PollInterval: wecomQRPollInterval, + PollTimeout: timeout, + Writer: os.Stdout, + } +} + +func applyWeComAuthResult(cfg *config.Config, botInfo wecomQRBotInfo) { + cfg.Channels.WeCom.Enabled = true + cfg.Channels.WeCom.BotID = botInfo.BotID + cfg.Channels.WeCom.SetSecret(botInfo.Secret) + if strings.TrimSpace(cfg.Channels.WeCom.WebSocketURL) == "" { + cfg.Channels.WeCom.WebSocketURL = wecomDefaultWebSocketURL + } +} + +func scanWeComQRCodeInteractive(ctx context.Context, opts wecomQRFlowOptions) (wecomQRBotInfo, error) { + opts = normalizeWeComQRFlowOptions(opts) + + fmt.Fprintln(opts.Writer, "Requesting WeCom QR code...") + + session, err := fetchWeComQRCode(ctx, opts) + if err != nil { + return wecomQRBotInfo{}, err + } + + fmt.Fprintln(opts.Writer) + fmt.Fprintln(opts.Writer, "=======================================================") + fmt.Fprintln(opts.Writer, "Please scan the following QR code with WeCom:") + fmt.Fprintln(opts.Writer, "=======================================================") + fmt.Fprintln(opts.Writer) + + qrterminal.GenerateWithConfig(session.AuthURL, qrterminal.Config{ + Level: qrterminal.L, + Writer: opts.Writer, + HalfBlocks: true, + }) + + pageURL, err := buildWeComQRCodePageURL(opts.QRCodePageURL, opts.SourceID, session.SCode) + if err != nil { + return wecomQRBotInfo{}, err + } + + fmt.Fprintln(opts.Writer) + fmt.Fprintf(opts.Writer, "QR Code Link: %s\n", pageURL) + fmt.Fprintln(opts.Writer) + fmt.Fprintln(opts.Writer, "Waiting for scan...") + + return pollWeComQRCodeResult(ctx, opts, session.SCode) +} + +func normalizeWeComQRFlowOptions(opts wecomQRFlowOptions) wecomQRFlowOptions { + if opts.HTTPClient == nil { + opts.HTTPClient = &http.Client{Timeout: wecomQRHTTPTimeout} + } + if strings.TrimSpace(opts.GenerateURL) == "" { + opts.GenerateURL = wecomQRGenerateEndpoint + } + if strings.TrimSpace(opts.QueryURL) == "" { + opts.QueryURL = wecomQRQueryEndpoint + } + if strings.TrimSpace(opts.QRCodePageURL) == "" { + opts.QRCodePageURL = wecomQRPageEndpoint + } + if strings.TrimSpace(opts.SourceID) == "" { + opts.SourceID = wecomQRSourceID + } + if opts.PollInterval <= 0 { + opts.PollInterval = wecomQRPollInterval + } + if opts.PollTimeout <= 0 { + opts.PollTimeout = wecomQRPollTimeout + } + if opts.Writer == nil { + opts.Writer = os.Stdout + } + + return opts +} + +func fetchWeComQRCode(ctx context.Context, opts wecomQRFlowOptions) (wecomQRSession, error) { + generateURL, err := buildWeComQRGenerateURL(opts.GenerateURL, opts.SourceID, wecomPlatformCode()) + if err != nil { + return wecomQRSession{}, err + } + + var resp wecomQRGenerateResponse + if err := doWeComJSONGet(ctx, opts.HTTPClient, generateURL, &resp); err != nil { + return wecomQRSession{}, fmt.Errorf("failed to get WeCom QR code: %w", err) + } + if resp.ErrCode != 0 { + return wecomQRSession{}, fmt.Errorf( + "failed to get WeCom QR code: errcode=%d errmsg=%s", + resp.ErrCode, + resp.ErrMsg, + ) + } + if resp.Data.SCode == "" || resp.Data.AuthURL == "" { + return wecomQRSession{}, fmt.Errorf("failed to get WeCom QR code: response missing scode or auth_url") + } + + return wecomQRSession{ + SCode: resp.Data.SCode, + AuthURL: resp.Data.AuthURL, + }, nil +} + +func pollWeComQRCodeResult(ctx context.Context, opts wecomQRFlowOptions, scode string) (wecomQRBotInfo, error) { + if strings.TrimSpace(scode) == "" { + return wecomQRBotInfo{}, fmt.Errorf("missing WeCom QR scode") + } + + timeoutCtx, cancel := context.WithTimeout(ctx, opts.PollTimeout) + defer cancel() + + var scannedPrinted bool + + for { + status, err := queryWeComQRCodeStatus(timeoutCtx, opts, scode) + if err != nil { + if errors.Is(err, context.DeadlineExceeded) || errors.Is(timeoutCtx.Err(), context.DeadlineExceeded) { + return wecomQRBotInfo{}, fmt.Errorf("WeCom QR scan timed out after %s", opts.PollTimeout) + } + return wecomQRBotInfo{}, err + } + + switch strings.ToLower(status.Data.Status) { + case "success": + if status.Data.BotInfo.BotID == "" || status.Data.BotInfo.Secret == "" { + return wecomQRBotInfo{}, fmt.Errorf("WeCom QR scan succeeded but bot credentials are missing") + } + return wecomQRBotInfo{ + BotID: status.Data.BotInfo.BotID, + Secret: status.Data.BotInfo.Secret, + }, nil + case "expired": + return wecomQRBotInfo{}, fmt.Errorf("WeCom QR code expired, please retry") + case "scaned", "scanned": + if !scannedPrinted { + fmt.Fprintln(opts.Writer, "QR code scanned. Confirm the login in WeCom.") + scannedPrinted = true + } + } + + select { + case <-timeoutCtx.Done(): + if errors.Is(timeoutCtx.Err(), context.DeadlineExceeded) { + return wecomQRBotInfo{}, fmt.Errorf("WeCom QR scan timed out after %s", opts.PollTimeout) + } + return wecomQRBotInfo{}, timeoutCtx.Err() + case <-time.After(opts.PollInterval): + } + } +} + +func queryWeComQRCodeStatus(ctx context.Context, opts wecomQRFlowOptions, scode string) (wecomQRQueryResponse, error) { + queryURL, err := buildWeComQRQueryURL(opts.QueryURL, scode) + if err != nil { + return wecomQRQueryResponse{}, err + } + + var resp wecomQRQueryResponse + if err := doWeComJSONGet(ctx, opts.HTTPClient, queryURL, &resp); err != nil { + return wecomQRQueryResponse{}, fmt.Errorf("failed to query WeCom QR result: %w", err) + } + if resp.ErrCode != 0 { + return wecomQRQueryResponse{}, fmt.Errorf( + "failed to query WeCom QR result: errcode=%d errmsg=%s", + resp.ErrCode, + resp.ErrMsg, + ) + } + + return resp, nil +} + +func buildWeComQRGenerateURL(baseURL, sourceID string, platformCode int) (string, error) { + u, err := url.Parse(baseURL) + if err != nil { + return "", fmt.Errorf("invalid WeCom QR generate URL: %w", err) + } + + query := u.Query() + query.Set("source", sourceID) + query.Set("sourceID", sourceID) + query.Set("plat", strconv.Itoa(platformCode)) + u.RawQuery = query.Encode() + + return u.String(), nil +} + +func buildWeComQRQueryURL(baseURL, scode string) (string, error) { + u, err := url.Parse(baseURL) + if err != nil { + return "", fmt.Errorf("invalid WeCom QR query URL: %w", err) + } + + query := u.Query() + query.Set("scode", scode) + u.RawQuery = query.Encode() + + return u.String(), nil +} + +func buildWeComQRCodePageURL(baseURL, sourceID, scode string) (string, error) { + u, err := url.Parse(baseURL) + if err != nil { + return "", fmt.Errorf("invalid WeCom QR page URL: %w", err) + } + + query := u.Query() + query.Set("source", sourceID) + query.Set("sourceID", sourceID) + query.Set("scode", scode) + u.RawQuery = query.Encode() + + return u.String(), nil +} + +func doWeComJSONGet(ctx context.Context, client *http.Client, targetURL string, out any) error { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, targetURL, nil) + if err != nil { + return err + } + + resp, err := client.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + body, readErr := io.ReadAll(io.LimitReader(resp.Body, 8192)) + if readErr != nil { + return fmt.Errorf("unexpected status %s", resp.Status) + } + return fmt.Errorf("unexpected status %s: %s", resp.Status, strings.TrimSpace(string(body))) + } + + if err := json.NewDecoder(resp.Body).Decode(out); err != nil { + return fmt.Errorf("decode JSON response: %w", err) + } + + return nil +} + +func wecomPlatformCode() int { + switch runtime.GOOS { + case "darwin": + return 1 + case "windows": + return 2 + case "linux": + return 3 + default: + return 0 + } +} diff --git a/cmd/picoclaw/internal/auth/wecom_test.go b/cmd/picoclaw/internal/auth/wecom_test.go new file mode 100644 index 000000000..c2a4624ae --- /dev/null +++ b/cmd/picoclaw/internal/auth/wecom_test.go @@ -0,0 +1,157 @@ +package auth + +import ( + "bytes" + "context" + "net/http" + "net/http/httptest" + "net/url" + "path/filepath" + "strconv" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/sipeed/picoclaw/cmd/picoclaw/internal" + "github.com/sipeed/picoclaw/pkg/config" +) + +func TestNewWeComCommand(t *testing.T) { + cmd := newWeComCommand() + + require.NotNil(t, cmd) + assert.Equal(t, "wecom", cmd.Use) + assert.Equal(t, "Scan a WeCom QR code and configure channels.wecom", cmd.Short) + assert.NotNil(t, cmd.Flags().Lookup("timeout")) +} + +func TestBuildWeComQRGenerateURL(t *testing.T) { + rawURL, err := buildWeComQRGenerateURL("https://example.com/ai/qc/generate", wecomQRSourceID, 3) + require.NoError(t, err) + + parsed, err := url.Parse(rawURL) + require.NoError(t, err) + + assert.Equal(t, wecomQRSourceID, parsed.Query().Get("source")) + assert.Equal(t, wecomQRSourceID, parsed.Query().Get("sourceID")) + assert.Equal(t, "3", parsed.Query().Get("plat")) +} + +func TestBuildWeComQRCodePageURL(t *testing.T) { + rawURL, err := buildWeComQRCodePageURL("https://example.com/ai/qc/gen", wecomQRSourceID, "scode-1") + require.NoError(t, err) + + parsed, err := url.Parse(rawURL) + require.NoError(t, err) + + assert.Equal(t, wecomQRSourceID, parsed.Query().Get("source")) + assert.Equal(t, wecomQRSourceID, parsed.Query().Get("sourceID")) + assert.Equal(t, "scode-1", parsed.Query().Get("scode")) +} + +func TestFetchWeComQRCode(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, "/generate", r.URL.Path) + assert.Equal(t, wecomQRSourceID, r.URL.Query().Get("source")) + assert.Equal(t, wecomQRSourceID, r.URL.Query().Get("sourceID")) + assert.Equal(t, strconv.Itoa(wecomPlatformCode()), r.URL.Query().Get("plat")) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"data":{"scode":"scode-1","auth_url":"https://example.com/qr"}}`)) + })) + defer server.Close() + + opts := normalizeWeComQRFlowOptions(wecomQRFlowOptions{ + HTTPClient: server.Client(), + GenerateURL: server.URL + "/generate", + Writer: bytes.NewBuffer(nil), + }) + + session, err := fetchWeComQRCode(context.Background(), opts) + require.NoError(t, err) + assert.Equal(t, "scode-1", session.SCode) + assert.Equal(t, "https://example.com/qr", session.AuthURL) +} + +func TestPollWeComQRCodeResult(t *testing.T) { + var calls atomic.Int32 + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + call := calls.Add(1) + assert.Equal(t, "/query", r.URL.Path) + assert.Equal(t, "scode-1", r.URL.Query().Get("scode")) + w.Header().Set("Content-Type", "application/json") + switch call { + case 1: + _, _ = w.Write([]byte(`{"data":{"status":"wait"}}`)) + case 2: + _, _ = w.Write([]byte(`{"data":{"status":"scaned"}}`)) + default: + _, _ = w.Write([]byte(`{"data":{"status":"success","bot_info":{"botid":"bot-1","secret":"secret-1"}}}`)) + } + })) + defer server.Close() + + var output bytes.Buffer + opts := normalizeWeComQRFlowOptions(wecomQRFlowOptions{ + HTTPClient: server.Client(), + QueryURL: server.URL + "/query", + PollInterval: time.Millisecond, + PollTimeout: time.Second, + Writer: &output, + }) + + botInfo, err := pollWeComQRCodeResult(context.Background(), opts, "scode-1") + require.NoError(t, err) + assert.Equal(t, "bot-1", botInfo.BotID) + assert.Equal(t, "secret-1", botInfo.Secret) + assert.Contains(t, output.String(), "QR code scanned. Confirm the login in WeCom.") +} + +func TestApplyWeComAuthResult(t *testing.T) { + cfg := config.DefaultConfig() + cfg.Channels.WeCom.WebSocketURL = "" + + applyWeComAuthResult(cfg, wecomQRBotInfo{ + BotID: "bot-1", + Secret: "secret-1", + }) + + assert.True(t, cfg.Channels.WeCom.Enabled) + assert.Equal(t, "bot-1", cfg.Channels.WeCom.BotID) + assert.Equal(t, "secret-1", cfg.Channels.WeCom.Secret()) + assert.Equal(t, wecomDefaultWebSocketURL, cfg.Channels.WeCom.WebSocketURL) +} + +func TestAuthWeComCmdWithScanner(t *testing.T) { + tmpDir := t.TempDir() + configPath := filepath.Join(tmpDir, "config.json") + + t.Setenv(config.EnvHome, tmpDir) + t.Setenv(config.EnvConfig, configPath) + + var output bytes.Buffer + err := authWeComCmdWithScanner( + context.Background(), + &output, + time.Second, + func(_ context.Context, opts wecomQRFlowOptions) (wecomQRBotInfo, error) { + assert.Equal(t, wecomQRSourceID, opts.SourceID) + return wecomQRBotInfo{ + BotID: "bot-1", + Secret: "secret-1", + }, nil + }, + ) + require.NoError(t, err) + + cfg, err := config.LoadConfig(internal.GetConfigPath()) + require.NoError(t, err) + assert.True(t, cfg.Channels.WeCom.Enabled) + assert.Equal(t, "bot-1", cfg.Channels.WeCom.BotID) + assert.Equal(t, "secret-1", cfg.Channels.WeCom.Secret()) + assert.Equal(t, wecomDefaultWebSocketURL, cfg.Channels.WeCom.WebSocketURL) + assert.Contains(t, output.String(), "WeCom connected.") +} diff --git a/config/config.example.json b/config/config.example.json index 88578701a..ff2969dcb 100644 --- a/config/config.example.json +++ b/config/config.example.json @@ -10,6 +10,7 @@ "max_tool_iterations": 20, "summarize_message_threshold": 20, "summarize_token_percent": 75, + "split_on_marker": false, "tool_feedback": { "enabled": false, "max_args_length": 300 @@ -129,6 +130,10 @@ "encrypt_key": "", "verification_token": "", "allow_from": [], + "placeholder": { + "enabled": true, + "text": ["Thinking...", "Processing...", "Typing..."] + }, "reasoning_channel_id": "", "random_reaction_emoji": [], "is_lark": false @@ -160,9 +165,11 @@ }, "placeholder": { "enabled": true, - "text": "Thinking... 💭" + "text": ["Thinking...", "Processing...", "Typing..."] }, - "reasoning_channel_id": "" + "reasoning_channel_id": "", + "crypto_database_path": "", + "crypto_passphrase": "YOUR_MATRIX_CRYPTO_PICKLE_KEY" }, "line": { "enabled": false, @@ -182,39 +189,13 @@ "reasoning_channel_id": "" }, "wecom": { - "_comment": "WeCom Bot - Easier setup, supports group chats", - "enabled": false, - "token": "YOUR_TOKEN", - "encoding_aes_key": "YOUR_43_CHAR_ENCODING_AES_KEY", - "webhook_url": "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=YOUR_KEY", - "webhook_path": "/webhook/wecom", - "allow_from": [], - "reply_timeout": 5, - "reasoning_channel_id": "" - }, - "wecom_app": { - "_comment": "WeCom App (自建应用) - More features, proactive messaging, private chat only.", - "enabled": false, - "corp_id": "YOUR_CORP_ID", - "corp_secret": "YOUR_CORP_SECRET", - "agent_id": 1000002, - "token": "YOUR_TOKEN", - "encoding_aes_key": "YOUR_43_CHAR_ENCODING_AES_KEY", - "webhook_path": "/webhook/wecom-app", - "allow_from": [], - "reply_timeout": 5, - "reasoning_channel_id": "" - }, - "wecom_aibot": { - "_comment": "WeCom AI Bot (智能机器人) - Official WeCom AI Bot integration, supports proactive messaging and private chats.", + "_comment": "WeCom AI Bot over WebSocket.", "enabled": false, "bot_id": "YOUR_BOT_ID", "secret": "YOUR_SECRET", - "token": "YOUR_TOKEN", - "encoding_aes_key": "YOUR_43_CHAR_ENCODING_AES_KEY", - "webhook_path": "/webhook/wecom-aibot", - "max_steps": 10, - "welcome_message": "Hello! I'm your AI assistant. How can I help you today?", + "websocket_url": "wss://openws.work.weixin.qq.com", + "send_thinking_message": true, + "allow_from": [], "reasoning_channel_id": "" }, "pico": { @@ -247,13 +228,8 @@ "nickserv_password": "", "sasl_user": "", "sasl_password": "", - "channels": [ - "#mychannel" - ], - "request_caps": [ - "server-time", - "message-tags" - ], + "channels": ["#mychannel"], + "request_caps": ["server-time", "message-tags"], "allow_from": [], "group_trigger": { "mention_only": true @@ -264,79 +240,6 @@ "reasoning_channel_id": "" } }, - "providers": { - "_comment": "DEPRECATED: Use model_list instead. This will be removed in a future version", - "anthropic": { - "api_key": "", - "api_base": "" - }, - "openai": { - "api_key": "", - "api_base": "", - "web_search": true - }, - "openrouter": { - "api_key": "sk-or-v1-xxx", - "api_base": "" - }, - "groq": { - "api_key": "gsk_xxx", - "api_base": "" - }, - "zhipu": { - "api_key": "YOUR_ZHIPU_API_KEY", - "api_base": "" - }, - "gemini": { - "api_key": "", - "api_base": "" - }, - "vllm": { - "api_key": "", - "api_base": "" - }, - "nvidia": { - "api_key": "nvapi-xxx", - "api_base": "", - "proxy": "http://127.0.0.1:7890" - }, - "moonshot": { - "api_key": "sk-xxx", - "api_base": "" - }, - "qwen": { - "api_key": "sk-xxx", - "api_base": "" - }, - "ollama": { - "api_key": "", - "api_base": "http://localhost:11434/v1" - }, - "cerebras": { - "api_key": "", - "api_base": "" - }, - "volcengine": { - "api_key": "", - "api_base": "" - }, - "mistral": { - "api_key": "", - "api_base": "https://api.mistral.ai/v1" - }, - "avian": { - "api_key": "", - "api_base": "https://api.avian.io/v1" - }, - "longcat": { - "api_key": "", - "api_base": "https://api.longcat.chat/openai" - }, - "modelscope": { - "api_key": "", - "api_base": "https://api-inference.modelscope.cn/v1" - } - }, "tools": { "allow_read_paths": null, "allow_write_paths": null, @@ -348,9 +251,7 @@ "brave": { "enabled": false, "api_key": "YOUR_BRAVE_API_KEY", - "api_keys": [ - "YOUR_BRAVE_API_KEY" - ], + "api_keys": ["YOUR_BRAVE_API_KEY"], "max_results": 5 }, "tavily": { @@ -366,9 +267,7 @@ "perplexity": { "enabled": false, "api_key": "pplx-xxx", - "api_keys": [ - "pplx-xxx" - ], + "api_keys": ["pplx-xxx"], "max_results": 5 }, "searxng": { @@ -383,6 +282,12 @@ "search_engine": "search_std", "max_results": 5 }, + "baidu_search": { + "enabled": false, + "api_key": "", + "base_url": "https://qianfan.baidubce.com/v2/ai_search/web_search", + "max_results": 10 + }, "fetch_limit_bytes": 10485760, "private_host_whitelist": [] }, @@ -411,19 +316,12 @@ "filesystem": { "enabled": false, "command": "npx", - "args": [ - "-y", - "@modelcontextprotocol/server-filesystem", - "/tmp" - ] + "args": ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"] }, "github": { "enabled": false, "command": "npx", - "args": [ - "-y", - "@modelcontextprotocol/server-github" - ], + "args": ["-y", "@modelcontextprotocol/server-github"], "env": { "GITHUB_PERSONAL_ACCESS_TOKEN": "YOUR_GITHUB_TOKEN" } @@ -431,10 +329,7 @@ "brave-search": { "enabled": false, "command": "npx", - "args": [ - "-y", - "@modelcontextprotocol/server-brave-search" - ], + "args": ["-y", "@modelcontextprotocol/server-brave-search"], "env": { "BRAVE_API_KEY": "YOUR_BRAVE_API_KEY" } @@ -451,10 +346,7 @@ "slack": { "enabled": false, "command": "npx", - "args": [ - "-y", - "@modelcontextprotocol/server-slack" - ], + "args": ["-y", "@modelcontextprotocol/server-slack"], "env": { "SLACK_BOT_TOKEN": "YOUR_SLACK_BOT_TOKEN", "SLACK_TEAM_ID": "YOUR_SLACK_TEAM_ID" diff --git a/docs/channels/matrix/README.md b/docs/channels/matrix/README.md index 2ed19245a..baded984e 100644 --- a/docs/channels/matrix/README.md +++ b/docs/channels/matrix/README.md @@ -22,10 +22,12 @@ Add this to `config.json`: }, "placeholder": { "enabled": true, - "text": "Thinking..." + "text": ["Thinking...", "Processing...", "Typing..."] }, "reasoning_channel_id": "", - "message_format": "richtext" + "message_format": "richtext", + "crypto_database_path": "", + "crypto_passphrase": "YOUR_MATRIX_CRYPTO_PICKLE_KEY" } } } @@ -43,9 +45,18 @@ Add this to `config.json`: | join_on_invite | bool | No | Auto-join invited rooms | | allow_from | []string | No | User whitelist (Matrix user IDs) | | group_trigger | object | No | Group trigger strategy (`mention_only` / `prefixes`) | -| placeholder | object | No | Placeholder message config | +| placeholder | object | No | Placeholder message config (see below) | | reasoning_channel_id | string | No | Target channel for reasoning output | | message_format | string | No | Output format: `"richtext"` (default) renders markdown as HTML; `"plain"` sends plain text only | +| crypto_database_path | string | No | Path to store the crypto database (uses workspace path `~/.picoclaw/workspace` if empty) | +| crypto_passphrase | string | No | Serialization key for encrypting session keys in the database; must remain unchanged once set | + +### Placeholder Config + +| Field | Type | Required | Description | +|---------|----------------|----------|-------------| +| enabled | bool | No | Enable placeholder messages (default: false) | +| text | string/[]string | No | Placeholder text(s). Can be a single string or array of strings. If multiple texts are provided, one is randomly selected at runtime. Default: "Thinking..." | ## 3. Currently Supported @@ -58,6 +69,7 @@ Add this to `config.json`: - Typing state (`m.typing`) - Placeholder message + final reply replacement - Auto-join invited rooms (can be disabled) +- End-to-end encryption (E2EE) support for encrypted messages ## 4. TODO diff --git a/docs/channels/matrix/README.zh.md b/docs/channels/matrix/README.zh.md index 8db3e4383..81afa550b 100644 --- a/docs/channels/matrix/README.zh.md +++ b/docs/channels/matrix/README.zh.md @@ -22,9 +22,12 @@ }, "placeholder": { "enabled": true, - "text": "Thinking... 💭" + "text": ["Thinking...", "Processing...", "Typing..."] }, - "reasoning_channel_id": "" + "reasoning_channel_id": "", + "message_format": "richtext", + "crypto_database_path": "", + "crypto_passphrase": "YOUR_MATRIX_CRYPTO_PICKLE_KEY" } } } @@ -45,6 +48,15 @@ | placeholder | object | 否 | 占位消息配置 | | reasoning_channel_id | string | 否 | 思维链输出目标通道 | | message_format | string | 否 | 消息格式：`richtext`（富文本）或 `plain`（纯文本） | +| crypto_database_path | string | 否 | 加密数据库存储路径（为空时使用工作空间路径 `~/.picoclaw/workspace`） | +| crypto_passphrase | string | 否 | 加密数据库中 session key 的序列化密钥；设置后不能更改 | + +### 占位消息配置 (Placeholder) + +| 字段 | 类型 | 必填 | 说明 | +|---------|-----------------|------|------| +| enabled | bool | 否 | 是否启用占位消息（默认：false） | +| text | string/[]string | 否 | 占位文本。可以是单个字符串或字符串数组。如果提供多个文本，运行时会随机选择一个。默认："Thinking..." | ## 3. 当前支持 @@ -56,6 +68,7 @@ - Typing 状态（`m.typing`） - 占位消息（`Thinking... 💭`）+ 最终回复替换 - 自动加入邀请房间（可关闭） +- 端对端加密（E2EE）消息支持 ## 4. TODO diff --git a/docs/channels/wecom/README.md b/docs/channels/wecom/README.md new file mode 100644 index 000000000..ecdfbc47b --- /dev/null +++ b/docs/channels/wecom/README.md @@ -0,0 +1,104 @@ +> Back to [README](../../../README.md) + +# WeCom + +PicoClaw now exposes WeCom as a single `channels.wecom` channel built on the official WeCom AI Bot WebSocket API. +This replaces the legacy `wecom`, `wecom_app`, and `wecom_aibot` split with one configuration model. + +## What This Channel Supports + +- Direct chat and group chat delivery +- Channel-side streaming replies over WeCom's AI Bot protocol +- Incoming text, voice, image, file, video, and mixed messages +- Outbound text and media replies (`image`, `file`, `voice`, `video`) +- QR-based CLI onboarding with `picoclaw auth wecom` +- Shared allowlist and `reasoning_channel_id` routing + +> No public webhook callback URL is required for this channel. PicoClaw opens an outbound WebSocket connection to WeCom. + +## Quick Start + +### Option 1: QR Login From CLI + +Run: + +```bash +picoclaw auth wecom +``` + +The command prints a QR code in the terminal, waits for confirmation in WeCom, and then writes the resulting +`bot_id` and `secret` into `channels.wecom`. + +Use `--timeout` if you want to wait longer: + +```bash +picoclaw auth wecom --timeout 10m +``` + +### Option 2: Configure Manually + +```json +{ + "channels": { + "wecom": { + "enabled": true, + "bot_id": "YOUR_BOT_ID", + "secret": "YOUR_SECRET", + "websocket_url": "wss://openws.work.weixin.qq.com", + "send_thinking_message": true, + "allow_from": [], + "reasoning_channel_id": "" + } + } +} +``` + +## Configuration + +| Field | Type | Required | Description | +| ----- | ---- | -------- | ----------- | +| `enabled` | bool | No | Enables the WeCom channel. | +| `bot_id` | string | Yes | WeCom AI Bot identifier. Required when the channel is enabled. | +| `secret` | string | Yes | WeCom AI Bot secret. Required when the channel is enabled. | +| `websocket_url` | string | No | WebSocket endpoint. Defaults to `wss://openws.work.weixin.qq.com`. | +| `send_thinking_message` | bool | No | Sends an initial `Processing...` chunk before the final streamed reply. Defaults to `true`. | +| `allow_from` | array | No | Sender allowlist. Empty means allow all senders. | +| `reasoning_channel_id` | string | No | Optional destination for reasoning/thinking output. | + +## Runtime Behavior + +- PicoClaw keeps the active WeCom turn so normal replies can continue the same stream when possible. +- If streaming is no longer available, replies fall back to active push delivery to the resolved chat route. +- Incoming media is downloaded into the media store before being handed to the agent. +- Outbound media is uploaded to WeCom in temporary chunks and then sent as a regular media message. + +## Migration Notes + +This branch removes the old multi-channel WeCom model. + +| Previous config | Now | +| --------------- | --- | +| `channels.wecom` webhook bot | Replace with `channels.wecom` using `bot_id` + `secret`. | +| `channels.wecom_app` | Remove it and use `channels.wecom`. | +| `channels.wecom_aibot` | Move the config to `channels.wecom`. | +| `token`, `encoding_aes_key`, `webhook_url`, `webhook_path` | No longer used by the WeCom channel. | +| `corp_id`, `corp_secret`, `agent_id` | No longer used by the WeCom channel. | +| `welcome_message`, `processing_message`, `max_steps` under WeCom | No longer part of the WeCom channel config. | + +## Troubleshooting + +### `picoclaw auth wecom` times out + +- Re-run with a larger `--timeout`. +- Make sure the QR code was confirmed inside WeCom, not only scanned. + +### WebSocket connection fails + +- Verify `bot_id` and `secret`. +- Confirm the host can reach `wss://openws.work.weixin.qq.com`. + +### Replies do not arrive + +- Check whether `allow_from` blocks the sender. +- Check launcher or startup validation for missing `channels.wecom.bot_id` / `channels.wecom.secret`. + diff --git a/docs/channels/wecom/README.zh.md b/docs/channels/wecom/README.zh.md new file mode 100644 index 000000000..6b4a5e495 --- /dev/null +++ b/docs/channels/wecom/README.zh.md @@ -0,0 +1,104 @@ +> 返回 [README](../../../README.zh.md) + +# 企业微信 + +PicoClaw 现在将企业微信统一为一个 `channels.wecom` 渠道，并基于企业微信官方 AI Bot WebSocket 协议实现。 +这取代了旧的 `wecom`、`wecom_app`、`wecom_aibot` 三套配置模型。 + +## 当前渠道能力 + +- 支持私聊和群聊 +- 支持企业微信侧流式回复 +- 支持接收文本、语音、图片、文件、视频和 mixed 消息 +- 支持发送文本与媒体消息（`image`、`file`、`voice`、`video`） +- 支持通过 `picoclaw auth wecom` 扫码写入配置 +- 支持统一白名单与 `reasoning_channel_id` + +> 这个渠道不再需要公网 webhook 回调地址。PicoClaw 会主动向企业微信发起 WebSocket 连接。 + +## 快速开始 + +### 方式 1：命令行扫码登录 + +运行： + +```bash +picoclaw auth wecom +``` + +该命令会在终端打印二维码，等待你在企业微信中确认，然后把生成的 `bot_id` 和 `secret` 写入 +`channels.wecom`。 + +如果需要更长等待时间，可以加 `--timeout`： + +```bash +picoclaw auth wecom --timeout 10m +``` + +### 方式 2：手动配置 + +```json +{ + "channels": { + "wecom": { + "enabled": true, + "bot_id": "YOUR_BOT_ID", + "secret": "YOUR_SECRET", + "websocket_url": "wss://openws.work.weixin.qq.com", + "send_thinking_message": true, + "allow_from": [], + "reasoning_channel_id": "" + } + } +} +``` + +## 配置字段 + +| 字段 | 类型 | 必填 | 说明 | +| ---- | ---- | ---- | ---- | +| `enabled` | bool | 否 | 是否启用企业微信渠道。 | +| `bot_id` | string | 是 | 企业微信 AI Bot 标识。渠道启用时必填。 | +| `secret` | string | 是 | 企业微信 AI Bot 密钥。渠道启用时必填。 | +| `websocket_url` | string | 否 | WebSocket 地址，默认 `wss://openws.work.weixin.qq.com`。 | +| `send_thinking_message` | bool | 否 | 是否在流式最终回复前先发送一段 `Processing...` 开场消息，默认 `true`。 | +| `allow_from` | array | 否 | 发送者白名单；空数组表示允许所有发送者。 | +| `reasoning_channel_id` | string | 否 | 可选的 reasoning/thinking 输出目标。 | + +## 运行时行为 + +- PicoClaw 会保留当前会话对应的企业微信 turn，优先继续同一个流式回复。 +- 如果流式上下文已经失效，回复会自动回退到主动推送消息。 +- 收到的媒体会先下载到 media store，再交给 Agent 处理。 +- 发出的媒体会先按分片上传到企业微信，再作为普通媒体消息发送。 + +## 迁移说明 + +这个分支移除了旧的多通道企业微信模型。 + +| 旧配置 | 现在怎么做 | +| ------ | ---------- | +| `channels.wecom` webhook 机器人 | 改为使用 `bot_id` + `secret` 的 `channels.wecom`。 | +| `channels.wecom_app` | 删除，统一迁移到 `channels.wecom`。 | +| `channels.wecom_aibot` | 配置迁移到 `channels.wecom`。 | +| `token`、`encoding_aes_key`、`webhook_url`、`webhook_path` | 企业微信渠道不再使用这些字段。 | +| `corp_id`、`corp_secret`、`agent_id` | 企业微信渠道不再使用这些字段。 | +| 企业微信下的 `welcome_message`、`processing_message`、`max_steps` | 不再属于企业微信渠道配置。 | + +## 常见问题 + +### `picoclaw auth wecom` 超时 + +- 用更大的 `--timeout` 重新执行。 +- 确认是在企业微信里完成了确认，而不只是扫描二维码。 + +### WebSocket 连接失败 + +- 检查 `bot_id` 和 `secret` 是否正确。 +- 确认运行环境可以访问 `wss://openws.work.weixin.qq.com`。 + +### 消息没有回到企业微信 + +- 检查 `allow_from` 是否拦截了发送者。 +- 检查启动日志或 launcher 校验，确认 `channels.wecom.bot_id` / `channels.wecom.secret` 已填写。 + diff --git a/docs/chat-apps.md b/docs/chat-apps.md index 4a78f465e..3d01994ff 100644 --- a/docs/chat-apps.md +++ b/docs/chat-apps.md @@ -6,7 +6,7 @@ Talk to your picoclaw through Telegram, Discord, WhatsApp, Matrix, QQ, DingTalk, LINE, WeCom, Feishu, Slack, IRC, OneBot, MaixCam, or Pico (native protocol) -> **Note**: All webhook-based channels (LINE, WeCom, etc.) are served on a single shared Gateway HTTP server (`gateway.host`:`gateway.port`, default `127.0.0.1:18790`). There are no per-channel ports to configure. Note: Feishu uses WebSocket/SDK mode and does not use the shared HTTP webhook server. +> **Note**: Channels that rely on HTTP callbacks share a single Gateway HTTP server (`gateway.host`:`gateway.port`, default `127.0.0.1:18790`). Socket/stream-based channels such as Feishu, DingTalk, and WeCom do not rely on the shared webhook server for inbound delivery. | Channel | Difficulty | Description | Documentation | | -------------------- | ------------------ | ----------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------- | @@ -19,7 +19,7 @@ Talk to your picoclaw through Telegram, Discord, WhatsApp, Matrix, QQ, DingTalk, | **QQ** | ⭐⭐ Medium | Official bot API, Chinese community | [Docs](channels/qq/README.md) | | **DingTalk** | ⭐⭐ Medium | Stream mode (no public IP needed), enterprise | [Docs](channels/dingtalk/README.md) | | **LINE** | ⭐⭐⭐ Advanced | HTTPS Webhook required | [Docs](channels/line/README.md) | -| **WeCom (企业微信)** | ⭐⭐⭐ Advanced | Group Bot (Webhook), custom App (API), AI Bot | [Bot](channels/wecom/wecom_bot/README.md) / [App](channels/wecom/wecom_app/README.md) / [AI Bot](channels/wecom/wecom_aibot/README.md) | +| **WeCom (企业微信)** | ⭐⭐⭐ Advanced | Official AI Bot over WebSocket, streaming + media | [Docs](channels/wecom/README.md) | | **Feishu (飞书)** | ⭐⭐⭐ Advanced | Enterprise collaboration, feature-rich | [Docs](channels/feishu/README.md) | | **IRC** | ⭐⭐ Medium | Server + TLS configuration | [Docs](#irc) | | **OneBot** | ⭐⭐ Medium | NapCat/Go-CQHTTP compatible, community ecosystem | [Docs](channels/onebot/README.md) | @@ -380,102 +380,34 @@ picoclaw gateway
WeCom (企业微信) -PicoClaw supports three types of WeCom integration: +PicoClaw now exposes WeCom as a single AI Bot channel over WebSocket. +No public webhook callback URL is required. -**Option 1: WeCom Bot (Bot)** - Easier setup, supports group chats -**Option 2: WeCom App (Custom App)** - More features, proactive messaging, private chat only -**Option 3: WeCom AI Bot (AI Bot)** - Official AI Bot, streaming replies, supports group & private chat +See [WeCom Configuration Guide](channels/wecom/README.md) for the full configuration reference and migration notes. -See [WeCom AI Bot Configuration Guide](channels/wecom/wecom_aibot/README.md) for detailed setup instructions. +**Quick Setup - Recommended** -**Quick Setup - WeCom Bot:** +**1. Authenticate** -**1. Create a bot** +```bash +picoclaw auth wecom +``` -* Go to WeCom Admin Console → Group Chat → Add Group Bot -* Copy the webhook URL (format: `https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=xxx`) +This command shows a QR code, waits for approval in WeCom, and writes `bot_id` + `secret` into `channels.wecom`. -**2. Configure** +**2. Configure manually if needed** ```json { "channels": { "wecom": { "enabled": true, - "token": "YOUR_TOKEN", - "encoding_aes_key": "YOUR_ENCODING_AES_KEY", - "webhook_url": "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=YOUR_KEY", - "webhook_path": "/webhook/wecom", - "allow_from": [] - } - } -} -``` - -> WeCom webhook is served on the shared Gateway server (`gateway.host`:`gateway.port`, default `127.0.0.1:18790`). - -**Quick Setup - WeCom App:** - -**1. Create an app** - -* Go to WeCom Admin Console → App Management → Create App -* Copy **AgentId** and **Secret** -* Go to "My Company" page, copy **CorpID** - -**2. Configure receive message** - -* In App details, click "Receive Message" → "Set API" -* Set URL to `http://your-server:18790/webhook/wecom-app` -* Generate **Token** and **EncodingAESKey** - -**3. Configure** - -```json -{ - "channels": { - "wecom_app": { - "enabled": true, - "corp_id": "wwxxxxxxxxxxxxxxxx", - "corp_secret": "YOUR_CORP_SECRET", - "agent_id": 1000002, - "token": "YOUR_TOKEN", - "encoding_aes_key": "YOUR_ENCODING_AES_KEY", - "webhook_path": "/webhook/wecom-app", - "allow_from": [] - } - } -} -``` - -**4. Run** - -```bash -picoclaw gateway -``` - -> **Note**: WeCom webhook callbacks are served on the Gateway port (default 18790). Use a reverse proxy for HTTPS. - -**Quick Setup - WeCom AI Bot:** - -**1. Create an AI Bot** - -* Go to WeCom Admin Console → App Management → AI Bot -* In the AI Bot settings, configure callback URL: `http://your-server:18790/webhook/wecom-aibot` -* Copy **Token** and click "Random Generate" for **EncodingAESKey** - -**2. Configure** - -```json -{ - "channels": { - "wecom_aibot": { - "enabled": true, - "token": "YOUR_TOKEN", - "encoding_aes_key": "YOUR_43_CHAR_ENCODING_AES_KEY", - "webhook_path": "/webhook/wecom-aibot", + "bot_id": "YOUR_BOT_ID", + "secret": "YOUR_SECRET", + "websocket_url": "wss://openws.work.weixin.qq.com", + "send_thinking_message": true, "allow_from": [], - "welcome_message": "Hello! How can I help you?", - "processing_message": "⏳ Processing, please wait. The results will be sent shortly." + "reasoning_channel_id": "" } } } @@ -487,7 +419,7 @@ picoclaw gateway picoclaw gateway ``` -> **Note**: WeCom AI Bot uses streaming pull protocol — no reply timeout concerns. Long tasks (>30 seconds) automatically switch to `response_url` push delivery. +> Legacy `wecom_app` and `wecom_aibot` entries are replaced by the unified `channels.wecom` config in this branch.
diff --git a/docs/configuration.md b/docs/configuration.md index 4e77300cf..9360d3897 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -31,6 +31,22 @@ PICOCLAW_HOME=/opt/picoclaw picoclaw agent PICOCLAW_HOME=/srv/picoclaw PICOCLAW_CONFIG=/srv/picoclaw/main.json picoclaw gateway ``` +### Gateway Log Level + +`gateway.log_level` controls Gateway log verbosity and is configurable in `config.json`. + +```json +{ + "gateway": { + "log_level": "fatal" + } +} +``` + +When omitted, the default is `fatal`. Supported values: `debug`, `info`, `warn`, `error`, `fatal`. + +You can also override this with the environment variable `PICOCLAW_LOG_LEVEL`. + ### Workspace Layout PicoClaw stores data in your configured workspace (default: `~/.picoclaw/workspace`): @@ -454,6 +470,70 @@ This design also enables **multi-agent support** with flexible provider selectio - **Load balancing**: Distribute requests across multiple endpoints - **Centralized configuration**: Manage all providers in one place +#### 🔒 Security Configuration (Recommended) + +PicoClaw supports separating sensitive data (API keys, tokens, secrets) from your main configuration by storing them in a `.security.yml` file. + +**Key Benefits:** +- **Security**: Sensitive data is never in your main config file +- **Easy sharing**: Share config.json without exposing API keys +- **Version control**: Add `.security.yml` to `.gitignore` +- **Flexible deployment**: Different environments can use different security files + +**Quick Setup:** + +1. Create `~/.picoclaw/.security.yml` with your API keys: +```yaml +model_list: + gpt-5.4: + api_keys: + - "sk-proj-your-actual-openai-key" + claude-sonnet-4.6: + api_keys: + - "sk-ant-your-actual-anthropic-key" +channels: + telegram: + token: "your-telegram-bot-token" +web: + brave: + api_keys: + - "BSAyour-brave-api-key" + glm_search: + api_key: "your-glm-search-api-key" +``` + +2. Set proper permissions: +```bash +chmod 600 ~/.picoclaw/.security.yml +``` + +3. Remove sensitive fields from `config.json` (recommended): +```json +{ + "model_list": [ + { + "model_name": "gpt-5.4", + "model": "openai/gpt-5.4" + // api_key loaded from .security.yml + } + ], + "channels": { + "telegram": { + "enabled": true" + // token loaded from .security.yml + } + } +} +``` + +**How it works:** +- Values from `.security.yml` are automatically mapped to config fields +- No special syntax needed — just omit sensitive fields from config.json +- If a field exists in both files, `.security.yml` value takes precedence +- You can mix direct values in config.json with security values + +For complete documentation, see [`security_configuration.md`](security_configuration.md). + #### All Supported Vendors | Vendor | `model` Prefix | Default API Base | Protocol | API Key | @@ -515,16 +595,20 @@ This design also enables **multi-agent support** with flexible provider selectio } ``` +> **Security Note**: You can remove `api_key` fields from your config and store them in `.security.yml` instead. See [Security Configuration](#-security-configuration-recommended) above for details. + #### Vendor-Specific Examples +> **Tip**: You can omit `api_key` fields and store them in `.security.yml` for better security. See [Security Configuration](#-security-configuration-recommended). +
OpenAI ```json { "model_name": "gpt-5.4", - "model": "openai/gpt-5.4", - "api_key": "sk-..." + "model": "openai/gpt-5.4" + // api_key: set in .security.yml } ``` @@ -536,8 +620,8 @@ This design also enables **multi-agent support** with flexible provider selectio ```json { "model_name": "ark-code-latest", - "model": "volcengine/ark-code-latest", - "api_key": "sk-..." + "model": "volcengine/ark-code-latest" + // api_key: set in .security.yml } ``` @@ -549,8 +633,8 @@ This design also enables **multi-agent support** with flexible provider selectio ```json { "model_name": "glm-4.7", - "model": "zhipu/glm-4.7", - "api_key": "your-key" + "model": "zhipu/glm-4.7" + // api_key: set in .security.yml } ``` @@ -562,8 +646,8 @@ This design also enables **multi-agent support** with flexible provider selectio ```json { "model_name": "deepseek-chat", - "model": "deepseek/deepseek-chat", - "api_key": "sk-..." + "model": "deepseek/deepseek-chat" + // api_key: set in .security.yml } ``` @@ -575,8 +659,8 @@ This design also enables **multi-agent support** with flexible provider selectio ```json { "model_name": "claude-sonnet-4.6", - "model": "anthropic/claude-sonnet-4.6", - "api_key": "sk-ant-your-key" + "model": "anthropic/claude-sonnet-4.6" + // api_key: set in .security.yml } ``` @@ -616,8 +700,8 @@ For direct Anthropic API access or custom endpoints that only support Anthropic' { "model_name": "my-custom-model", "model": "openai/custom-model", - "api_base": "https://my-proxy.com/v1", - "api_key": "sk-..." + "api_base": "https://my-proxy.com/v1" + // api_key: set in .security.yml } ``` @@ -629,6 +713,33 @@ PicoClaw strips only the outer `litellm/` prefix before sending the request, so Configure multiple endpoints for the same model name — PicoClaw will automatically round-robin between them: +**Option 1: Multiple API Keys in .security.yml (Recommended)** + +```yaml +# .security.yml +model_list: + gpt-5.4: + api_keys: + - "sk-proj-key-1" + - "sk-proj-key-2" +``` + +```json +// config.json +{ + "model_list": [ + { + "model_name": "gpt-5.4", + "model": "openai/gpt-5.4", + "api_base": "https://api.openai.com/v1" + // api_keys loaded from .security.yml + } + ] +} +``` + +**Option 2: Multiple Model Entries** + ```json { "model_list": [ @@ -685,6 +796,8 @@ This keeps the runtime lightweight while making new OpenAI-compatible backends m } ``` +> **Note**: The `providers` format is deprecated. Use the new `model_list` format with `.security.yml` for better security. +
@@ -701,18 +814,10 @@ This keeps the runtime lightweight while making new OpenAI-compatible backends m "dm_scope": "per-channel-peer", "backlog_limit": 20 }, - "providers": { - "openrouter": { - "api_key": "sk-or-v1-xxx" - }, - "groq": { - "api_key": "gsk_xxx" - } - }, "channels": { "telegram": { - "enabled": true, - "token": "123456:ABC...", + "enabled": true" + // token: set in .security.yml "allow_from": ["123456789"] } }, @@ -731,6 +836,8 @@ This keeps the runtime lightweight while making new OpenAI-compatible backends m } ``` +> **Note**: Sensitive fields (`api_key`, `token`, etc.) can be omitted and stored in `.security.yml` for better security. +
### Scheduled Tasks / Reminders @@ -754,6 +861,7 @@ Scheduled tasks persist across restarts and are stored in `~/.picoclaw/workspace | Topic | Description | | ----- | ----------- | +| [Security Configuration](security_configuration.md) | Store API keys and secrets in separate `.security.yml` file | | [Sensitive Data Filtering](sensitive_data_filtering.md) | Filter API keys and tokens from tool results before sending to LLM | | [Hook System](hooks/README.md) | Event-driven hooks: observers, interceptors, approval hooks | | [Steering](steering.md) | Inject messages into a running agent loop between tool calls | diff --git a/docs/credential_encryption.md b/docs/credential_encryption.md index dde8c782c..de3b70e09 100644 --- a/docs/credential_encryption.md +++ b/docs/credential_encryption.md @@ -31,7 +31,7 @@ enc://AAAA...base64... { "model_name": "gpt-4o", "model": "openai/gpt-4o", - "api_key": "enc://AAAA...base64...", + // "api_key": "enc://AAAA...base64..." move to .security.yml "api_base": "https://api.openai.com/v1" } ] diff --git a/docs/docker.md b/docs/docker.md index f868d4a42..a00dfbe9f 100644 --- a/docs/docker.md +++ b/docs/docker.md @@ -26,6 +26,9 @@ docker compose -f docker/docker-compose.yml --profile gateway up -d > [!TIP] > **Docker Users**: By default, the Gateway listens on `127.0.0.1` which is not accessible from the host. If you need to access the health endpoints or expose ports, set `PICOCLAW_GATEWAY_HOST=0.0.0.0` in your environment or update `config.json`. +> [!NOTE] +> The `gateway` profile only serves the webhook handlers (including Pico when enabled) and health endpoints on the gateway port, so it does not expose generic REST chat endpoints such as `/chat` or `/a2a`. Launcher mode adds the browser UI plus `/api/pico/token` and a `/pico/ws` proxy on the launcher port, but `/pico/ws` is also available directly on the gateway whenever the Pico channel is enabled. + ```bash # 5. Check logs docker compose -f docker/docker-compose.yml logs -f picoclaw-gateway diff --git a/docs/providers.md b/docs/providers.md index 3a740d3b8..c6d442a3b 100644 --- a/docs/providers.md +++ b/docs/providers.md @@ -11,6 +11,7 @@ | ------------ | --------------------------------------- | ------------------------------------------------------------ | | `gemini` | LLM (Gemini direct) | [aistudio.google.com](https://aistudio.google.com) | | `zhipu` | LLM (Zhipu direct) | [bigmodel.cn](https://bigmodel.cn) | +| `zai-coding` | LLM (Z.AI Coding Plan) | [z.ai](https://z.ai/manage-apikey/apikey-list) | | `volcengine` | LLM(Volcengine direct) | [volcengine.com](https://www.volcengine.com/activity/codingplan?utm_campaign=PicoClaw&utm_content=PicoClaw&utm_medium=devrel&utm_source=OWO&utm_term=PicoClaw) | | `openrouter` | LLM (recommended, access to all models) | [openrouter.ai](https://openrouter.ai) | | `anthropic` | LLM (Claude direct) | [console.anthropic.com](https://console.anthropic.com) | @@ -27,6 +28,7 @@ | `mistral` | LLM (Mistral direct) | [console.mistral.ai](https://console.mistral.ai) | | `longcat` | LLM (Longcat direct) | [longcat.ai](https://longcat.ai) | | `modelscope` | LLM (ModelScope direct) | [modelscope.cn](https://modelscope.cn) | +| `mimo` | LLM (Xiaomi MiMo direct) | [platform.xiaomimimo.com](https://platform.xiaomimimo.com) | ### Model Configuration (model_list) @@ -46,6 +48,7 @@ This design also enables **multi-agent support** with flexible provider selectio | **OpenAI** | `openai/` | `https://api.openai.com/v1` | OpenAI | [Get Key](https://platform.openai.com) | | **Anthropic** | `anthropic/` | `https://api.anthropic.com/v1` | Anthropic | [Get Key](https://console.anthropic.com) | | **智谱 AI (GLM)** | `zhipu/` | `https://open.bigmodel.cn/api/paas/v4` | OpenAI | [Get Key](https://open.bigmodel.cn/usercenter/proj-mgmt/apikeys) | +| **Z.AI Coding Plan** | `openai/` | `https://api.z.ai/api/coding/paas/v4` | OpenAI | [Get Key](https://z.ai/manage-apikey/apikey-list) | | **DeepSeek** | `deepseek/` | `https://api.deepseek.com/v1` | OpenAI | [Get Key](https://platform.deepseek.com) | | **Google Gemini** | `gemini/` | `https://generativelanguage.googleapis.com/v1beta` | OpenAI | [Get Key](https://aistudio.google.com/api-keys) | | **Groq** | `groq/` | `https://api.groq.com/openai/v1` | OpenAI | [Get Key](https://console.groq.com) | @@ -63,6 +66,7 @@ This design also enables **multi-agent support** with flexible provider selectio | **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) | +| **Xiaomi MiMo** | `mimo/` | `https://api.xiaomimimo.com/v1` | OpenAI | [Get Key](https://platform.xiaomimimo.com) | | **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 | - | @@ -160,6 +164,17 @@ If `voice.model_name` is not configured, PicoClaw will continue to fall back to } ``` +**Z.AI Coding Plan (GLM)** +> Z.AI and 智谱 AI are two brands of the same provider. For the Z.AI Coding Plan use the `openai` model key and the api base as follows, rather than the zhipu config +```json +{ + "model_name": "glm-4.7", + "model": "openai/glm-4.7", + "api_key": "your-z.ai-key" + "api_base": "https://api.z.ai/api/coding/paas/v4" +} +``` + **DeepSeek** ```json @@ -236,6 +251,21 @@ For direct Anthropic API access or custom endpoints that only support Anthropic' PicoClaw strips only the outer `litellm/` prefix before sending the request, so proxy aliases like `litellm/lite-gpt4` send `lite-gpt4`, while `litellm/openai/gpt-4o` sends `openai/gpt-4o`. +**Z.AI Coding Plan** + +If the standard Zhipu endpoint (`https://open.bigmodel.cn/api/paas/v4`) returns 429 (code 1113: insufficient balance), try using the Z.AI Coding Plan endpoint instead: + +```json +{ + "model_name": "glm-4.7", + "model": "openai/glm-4.7", + "api_key": "your-zhipu-api-key", + "api_base": "https://api.z.ai/api/coding/paas/v4" +} +``` + +**Note:** The Z.AI Coding Plan endpoint and standard Zhipu endpoint use the same API key format but have separate billing. If you encounter 429 errors with the standard Zhipu endpoint, the Z.AI Coding Plan endpoint may have available balance. + #### Load Balancing Configure multiple endpoints for the same model name—PicoClaw will automatically round-robin between them: @@ -259,6 +289,45 @@ Configure multiple endpoints for the same model name—PicoClaw will automatical } ``` +#### Automatic Model Failover (Cascade) + +PicoClaw already supports automatic failover when you configure `primary` + `fallbacks` in the agent model settings. +The runtime fallback chain retries the next candidate for retriable failures such as HTTP `429`, quota/rate-limit errors, and timeout errors. +It also applies cooldown tracking per candidate to avoid immediately retrying a recently failed target. + +```json +{ + "model_list": [ + { + "model_name": "qwen-main", + "model": "openai/qwen3.5:cloud", + "api_base": "https://api.example.com/v1", + "api_key": "sk-main" + }, + { + "model_name": "deepseek-backup", + "model": "deepseek/deepseek-chat", + "api_key": "sk-backup-1" + }, + { + "model_name": "gemini-backup", + "model": "gemini/gemini-2.5-flash", + "api_key": "sk-backup-2" + } + ], + "agents": { + "defaults": { + "model": { + "primary": "qwen-main", + "fallbacks": ["deepseek-backup", "gemini-backup"] + } + } + } +} +``` + +If you use key-level failover for the same model, PicoClaw can chain through additional key-backed candidates before moving to cross-model backups. + #### Migration from Legacy `providers` Config The old `providers` configuration is **deprecated** but still supported for backward compatibility. diff --git a/docs/security_configuration.md b/docs/security_configuration.md new file mode 100644 index 000000000..f4fe0e304 --- /dev/null +++ b/docs/security_configuration.md @@ -0,0 +1,644 @@ +# Security Configuration + +## Overview + +PicoClaw supports separating sensitive data (API keys, tokens, secrets, passwords) from the main configuration by storing them in a `.security.yml` file. This improves security by: + +1. **Separation of concerns**: Configuration settings and secrets are in separate files +2. **Easier sharing**: The main config can be shared without exposing sensitive data +3. **Better version control**: `.security.yml` should be added to `.gitignore` +4. **Flexible deployment**: Different environments can use different security files + +## File Structure + +``` +~/.picoclaw/ +├── config.json # Main configuration (safe to share) +└── .security.yml # Security data (never share) +``` + +## How It Works + +The security configuration works through **direct field mapping**, NOT through `ref:` string references. The system automatically loads values from `.security.yml` and applies them to the corresponding fields in `config.json`. + +### Key Points: + +- Values in `.security.yml` are automatically mapped to corresponding fields in the config +- The mapping is based on field names and structure, not on reference strings +- If a value exists in `.security.yml`, it **overrides** the value in `config.json` +- You can omit sensitive fields from `config.json` entirely (recommended) + +## Security Configuration Structure + +### Complete Example: .security.yml + +```yaml +# Model API Keys +# All models MUST use `api_keys` (plural) array format +# Even a single key must be provided as an array with one element +model_list: + gpt-5.4: + api_keys: + - "sk-proj-your-actual-openai-key-1" + - "sk-proj-your-actual-openai-key-2" # Optional: Multiple keys for failover + claude-sonnet-4.6: + api_keys: + - "sk-ant-your-actual-anthropic-key" # Single key in array format + +# Channel Tokens +channels: + telegram: + token: "your-telegram-bot-token" + feishu: + app_secret: "your-feishu-app-secret" + encrypt_key: "your-feishu-encrypt-key" + verification_token: "your-feishu-verification-token" + discord: + token: "your-discord-bot-token" + weixin: + token: "your-weixin-token" + qq: + app_secret: "your-qq-app-secret" + dingtalk: + client_secret: "your-dingtalk-client-secret" + slack: + bot_token: "your-slack-bot-token" + app_token: "your-slack-app-token" + matrix: + access_token: "your-matrix-access-token" + line: + channel_secret: "your-line-channel-secret" + channel_access_token: "your-line-channel-access-token" + onebot: + access_token: "your-onebot-access-token" + wecom: + token: "your-wecom-token" + encoding_aes_key: "your-wecom-encoding-aes-key" + wecom_app: + corp_secret: "your-wecom-app-corp-secret" + token: "your-wecom-app-token" + encoding_aes_key: "your-wecom-app-encoding-aes-key" + wecom_aibot: + secret: "your-wecom-aibot-secret" + token: "your-wecom-aibot-token" + encoding_aes_key: "your-wecom-aibot-encoding-aes-key" + pico: + token: "your-pico-token" + irc: + password: "your-irc-password" + nickserv_password: "your-irc-nickserv-password" + sasl_password: "your-irc-sasl-password" + +# Web Tool API Keys +web: + brave: + api_keys: + - "BSAyour-brave-api-key-1" + - "BSAyour-brave-api-key-2" # Optional: Multiple keys for failover + tavily: + api_keys: + - "tvly-your-tavily-api-key" # Single key in array format + perplexity: + api_keys: + - "pplx-your-perplexity-api-key" # Single key in array format + glm_search: + api_key: "your-glm-search-api-key" # GLMSearch uses single key format (not array) + baidu_search: + api_key: "your-baidu-search-api-key" + +# Skills Registry Tokens +skills: + github: + token: "your-github-token" + clawhub: + auth_token: "your-clawhub-auth-token" +``` + +## Usage + +### Step 1: Create .security.yml + +Create or copy the security file: +```bash +cp security.example.yml ~/.picoclaw/.security.yml +``` + +### Step 2: Fill in your actual values + +Edit `~/.picoclaw/.security.yml` and replace placeholder values with your actual API keys and tokens. + +### Step 3: Set proper permissions + +```bash +chmod 600 ~/.picoclaw/.security.yml +``` + +### Step 4: Simplify config.json (Recommended) + +You can now remove sensitive fields from `config.json` since they're loaded from `.security.yml`: + +**Before:** +```json +{ + "model_list": [ + { + "model_name": "gpt-5.4", + "model": "openai/gpt-5.4", + "api_base": "https://api.openai.com/v1", + "api_key": "sk-your-actual-api-key-here" + } + ], + "channels": { + "telegram": { + "enabled": true, + "token": "1234567890:ABCdefGHIjklMNOpqrsTUVwxyz" + } + } +} +``` + +**After:** +```json +{ + "model_list": [ + { + "model_name": "gpt-5.4", + "model": "openai/gpt-5.4", + "api_base": "https://api.openai.com/v1" + // api_key is now loaded from .security.yml + } + ], + "channels": { + "telegram": { + "enabled": true" + // token is now loaded from .security.yml + } + } +} +``` + +### Step 5: Verify + +Restart PicoClaw and verify it loads correctly: +```bash +picoclaw --version +``` + +## Field Mapping Rules + +### Models + +**In .security.yml:** +```yaml +model_list: + : + api_keys: + - "key-1" + - "key-2" +``` + +**Mapping:** +- Field `api_keys` (array) maps to the model's API keys +- The `` must match the `model_name` field in `config.json` +- Supports indexed names (e.g., "gpt-5.4:0") - the system will also try the base name ("gpt-5.4") + +### Channels + +Each channel maps its fields directly: + +**In .security.yml:** +```yaml +channels: + telegram: + token: "value" + feishu: + app_secret: "value" + encrypt_key: "value" + verification_token: "value" + discord: + token: "value" +``` + +**Mapping:** +- `channels.telegram.token` → `config.channels.telegram.token` +- `channels.feishu.app_secret` → `config.channels.feishu.app_secret` +- etc. + +### Web Tools + +**Brave, Tavily, Perplexity:** +```yaml +web: + brave: + api_keys: + - "key-1" + - "key-2" +``` +- Use `api_keys` (plural) array format + +**GLMSearch:** +```yaml +web: + glm_search: + api_key: "single-key-here" +``` +- Use `api_key` (singular) single string format + +**BaiduSearch:** +```yaml +web: + baidu_search: + api_key: "your-key" +``` +- Use `api_key` (singular) single string format + +### Skills + +**In .security.yml:** +```yaml +skills: + github: + token: "value" + clawhub: + auth_token: "value" +``` + +## API Key Formats + +### Models - Single key + +Use array format with one element: +```yaml +model_list: + gpt-5.4: + api_keys: + - "sk-your-key" +``` + +### Models - Multiple keys (Load Balancing & Failover) + +Use array format with multiple elements: +```yaml +model_list: + gpt-5.4: + api_keys: + - "sk-your-key-1" + - "sk-your-key-2" + - "sk-your-key-3" +``` + +**Benefits:** +- **Load balancing**: Requests are distributed across multiple keys +- **Failover**: Automatic switching to another key if one fails +- **Rate limit management**: Distribute usage across multiple keys +- **High availability**: Reduce downtime during API provider issues + +### Web Tools (Brave/Tavily/Perplexity) - Single key + +```yaml +web: + brave: + api_keys: + - "BSA-your-key" +``` + +### Web Tools (Brave/Tavily/Perplexity) - Multiple keys + +```yaml +web: + brave: + api_keys: + - "BSA-key-1" + - "BSA-key-2" +``` + +### Web Tool (GLMSearch/BaiduSearch) - Single key only + +```yaml +web: + glm_search: + api_key: "your-glm-key" # Single string (NOT array) + baidu_search: + api_key: "your-baidu-key" # Single string (NOT array) +``` + +## Model Name Matching + +The system supports intelligent model name matching in `.security.yml`: + +### Example 1: Exact Match + +**config.json:** +```json +{ + "model_name": "gpt-5.4:0" +} +``` + +**.security.yml (exact match with index):** +```yaml +model_list: + gpt-5.4:0: + api_keys: ["key-1"] +``` + +### Example 2: Base Name Match + +**config.json:** +```json +{ + "model_name": "gpt-5.4:0" +} +``` + +**.security.yml (base name without index):** +```yaml +model_list: + gpt-5.4: + api_keys: ["key-1", "key-2"] +``` + +Both methods work. The base name match allows you to use simpler keys in `.security.yml` even when your config uses indexed model names for load balancing. + +## Backward Compatibility + +The system maintains full backward compatibility: + +1. **Direct values**: You can still use direct values in `config.json` (not recommended for production) +2. **Mixed usage**: You can have some fields in `.security.yml` and others in `config.json` +3. **Optional security file**: If `.security.yml` doesn't exist, the system will only use values from `config.json` +4. **Override behavior**: If a field exists in both files, `.security.yml` value takes precedence + +## Environment Variables + +You can override any security value using environment variables: + +**For models:** +```bash +export PICOCLAW_CHANNELS_TELEGRAM_TOKEN="token-from-env" +``` + +**For channels:** +```bash +export PICOCLAW_CHANNELS_TELEGRAM_TOKEN="token-from-env" +export PICOCLAW_CHANNELS_FEISHU_APP_SECRET="secret-from-env" +``` + +**For web tools:** +```bash +export PICOCLAW_TOOLS_WEB_BRAVE_API_KEY="key-from-env" +export PICOCLAW_TOOLS_WEB_BAIDU_API_KEY="baidu-key-from-env" +``` + +Environment variables have the highest priority and will override both `config.json` and `.security.yml` values. + +The pattern is: `PICOCLAW_
__` with underscores separating path segments and converted to uppercase. + +## Security Best Practices + +1. **Never commit `.security.yml`** to version control +2. **Add to .gitignore**: Ensure `.security.yml` is in your `.gitignore` file +3. **Set file permissions**: `chmod 600 ~/.picoclaw/.security.yml` +4. **Use different keys** for different environments (dev, staging, production) +5. **Rotate keys regularly** and update `.security.yml` +6. **Backup securely**: Encrypt backups containing `.security.yml` +7. **Review access**: Ensure only authorized users have read access to the file + +## API + +### loadSecurityConfig + +```go +func loadSecurityConfig(securityPath string) (*SecurityConfig, error) +``` + +Loads the security configuration from `.security.yml`. Returns an empty `SecurityConfig` if the file doesn't exist. + +### saveSecurityConfig + +```go +func saveSecurityConfig(securityPath string, sec *SecurityConfig) error +``` + +Saves the security configuration to `.security.yml` with `0o600` permissions. + +### applySecurityConfig + +```go +func applySecurityConfig(cfg *Config, sec *SecurityConfig) error +``` + +Applies security configuration to the main config by copying values from `.security.yml` to the corresponding fields in the config. + +### securityPath + +```go +func securityPath(configPath string) string +``` + +Returns the path to `.security.yml` relative to the config file. + +## Example: Complete Configuration + +### config.json + +```json +{ + "version": 1, + "agents": { + "defaults": { + "workspace": "~/picoclaw-workspace", + "model_name": "gpt-5.4" + } + }, + "model_list": [ + { + "model_name": "gpt-5.4", + "model": "openai/gpt-5.4", + "api_base": "https://api.openai.com/v1" + }, + { + "model_name": "claude-sonnet-4.6", + "model": "anthropic/claude-sonnet-4.6", + "api_base": "https://api.anthropic.com/v1" + } + ], + "channels": { + "telegram": { + "enabled": true + } + }, + "tools": { + "web": { + "brave": { + "enabled": true + } + } + } +} +``` + +### .security.yml + +```yaml +model_list: + gpt-5.4: + api_keys: + - "sk-proj-actual-openai-key-1" + - "sk-proj-actual-openai-key-2" + claude-sonnet-4.6: + api_keys: + - "sk-ant-actual-anthropic-key" + +channels: + telegram: + token: "1234567890:ABCdefGHIjklMNOpqrsTUVwxyz" + +web: + brave: + api_keys: + - "BSAactualbravekey-1" + - "BSAactualbravekey-2" + tavily: + api_keys: + - "tvly-your-tavily-key" + glm_search: + api_key: "your-glm-key" + baidu_search: + api_key: "your-baidu-key" +``` + +## Testing + +Run the security configuration tests: + +```bash +go test ./pkg/config -run TestSecurityConfig +``` + +## Troubleshooting + +### Error: "failed to load security config" + +- Verify `.security.yml` exists in the same directory as `config.json` +- Check the YAML syntax is valid (use a YAML validator) +- Ensure file permissions allow reading + +### Error: "model security entry not found" + +- Ensure the model name in `config.json` matches exactly in `.security.yml` +- Check that the `model_list` section exists in `.security.yml` +- For models with indexed names (e.g., "gpt-5.4:0"), ensure the exact name is used or check the base name without index +- Verify the YAML structure is correct (proper indentation) + +### Multiple API Keys Not Working + +- Ensure you're using `api_keys` (plural) in `.security.yml` for models and web tools (except GLMSearch/BaiduSearch) +- Check that the array format is correct in YAML (proper indentation with dashes) +- Remember: Models, Brave, Tavily, Perplexity MUST use `api_keys` (array format) +- GLMSearch and BaiduSearch MUST use `api_key` (single string format) + +### Load Balancing/Failover Issues + +- Verify all API keys in the `api_keys` array are valid +- Check that all keys have the same rate limits and permissions +- Monitor logs to see which keys are being used and failing +- Ensure the `api_keys` array is properly formatted in YAML + +### Keys Not Being Applied + +- Check that `.security.yml` is in the same directory as `config.json` +- Verify the file permissions allow reading (`chmod 600 ~/.picoclaw/.security.yml`) +- Ensure the YAML structure matches the expected format +- Check for typos in field names (case-sensitive) +- Verify the model/channel names match exactly (case-sensitive) + +## Migration Guide + +### Step 1: Backup your config + +```bash +cp ~/.picoclaw/config.json ~/.picoclaw/config.json.backup +``` + +### Step 2: Create .security.yml + +```bash +cp security.example.yml ~/.picoclaw/.security.yml +``` + +### Step 3: Fill in your API keys + +Edit `~/.picoclaw/.security.yml` and replace placeholder values with your actual keys. + +### Step 4: Remove sensitive fields from config.json + +Remove or comment out sensitive fields from `config.json`: +- `api_key` fields from `model_list` entries +- `token` fields from `channels` +- `api_key` fields from `tools.web` +- `token`/`auth_token` fields from `tools.skills` + +### Step 5: Set proper permissions + +```bash +chmod 600 ~/.picoclaw/.security.yml +``` + +### Step 6: Test + +```bash +picoclaw --version +``` + +### Step 7: Verify functionality + +Test your models and channels to ensure everything works correctly. + +### Step 8: Clean up (optional) + +If everything works, you can delete the backup: +```bash +rm ~/.picoclaw/config.json.backup +``` + +## Advanced: Encrypted API Keys + +PicoClaw supports encrypting API keys in the security file for additional protection. + +### Setup + +1. Set a passphrase via environment variable: +```bash +export PICOCLAW_CREDENTIAL_PASSPHRASE="your-secure-passphrase" +``` + +2. When saving config, API keys will be encrypted automatically: +```go +SaveConfig(path, config) +``` + +### Encrypted Format + +Encrypted keys are stored as: +```yaml +model_list: + gpt-5.4: + api_keys: + - "enc://encrypted-base64-string" +``` + +The system automatically decrypts keys at runtime when loading the configuration. + +### Benefits + +- Additional layer of security +- Keys are encrypted at rest +- Passphrase can be managed separately from the config file + +### Important Notes + +- Always backup your passphrase securely +- If you lose the passphrase, you'll lose access to encrypted keys +- Use a strong, unique passphrase +- Never commit the passphrase to version control diff --git a/docs/tools_configuration.md b/docs/tools_configuration.md index b5907b991..314cbd38f 100644 --- a/docs/tools_configuration.md +++ b/docs/tools_configuration.md @@ -70,12 +70,12 @@ General settings for fetching and processing webpage content. Baidu Search uses the [Qianfan AI Search API](https://cloud.baidu.com/doc/qianfan-api/s/Wmbq4z7e5), which is AI-powered and optimized for Chinese-language queries. -| Config | Type | Default | Description | -|---------------|--------|------------------------------------------------------------------|---------------------------| -| `enabled` | bool | false | Enable Baidu Search | -| `api_key` | string | - | Qianfan API key | -| `base_url` | string | `https://qianfan.baidubce.com/v2/ai_search/web_search` | Baidu Search API URL | -| `max_results` | int | 10 | Maximum number of results | +| Config | Type | Default | Description | +|---------------|--------|--------------------------------------------------------|---------------------------| +| `enabled` | bool | false | Enable Baidu Search | +| `api_key` | string | - | Qianfan API key | +| `base_url` | string | `https://qianfan.baidubce.com/v2/ai_search/web_search` | Baidu Search API URL | +| `max_results` | int | 5 | Maximum number of results | ```json { @@ -107,25 +107,25 @@ Baidu Search uses the [Qianfan AI Search API](https://cloud.baidu.com/doc/qianfa | `enabled` | bool | false | Enable Tavily search | | `api_key` | string | - | Tavily API key | | `base_url` | string | - | Custom Tavily API base URL | -| `max_results` | int | 0 | Maximum number of results (0 = default) | +| `max_results` | int | 5 | Maximum number of results | ### SearXNG -| Config | Type | Default | Description | -|---------------|--------|--------------------------|---------------------------| -| `enabled` | bool | false | Enable SearXNG search | -| `base_url` | string | `http://localhost:8888` | SearXNG instance URL | -| `max_results` | int | 5 | Maximum number of results | +| Config | Type | Default | Description | +|---------------|--------|-------------------------|---------------------------| +| `enabled` | bool | false | Enable SearXNG search | +| `base_url` | string | `http://localhost:8888` | SearXNG instance URL | +| `max_results` | int | 5 | Maximum number of results | ### GLM Search -| Config | Type | Default | Description | -|-----------------|--------|------------------------------------------------------|---------------------------| -| `enabled` | bool | false | Enable GLM Search | -| `api_key` | string | - | GLM API key | -| `base_url` | string | `https://open.bigmodel.cn/api/paas/v4/web_search` | GLM Search API URL | -| `search_engine` | string | `search_std` | Search engine type | -| `max_results` | int | 5 | Maximum number of results | +| Config | Type | Default | Description | +|-----------------|--------|---------------------------------------------------|---------------------------| +| `enabled` | bool | false | Enable GLM Search | +| `api_key` | string | - | GLM API key | +| `base_url` | string | `https://open.bigmodel.cn/api/paas/v4/web_search` | GLM Search API URL | +| `search_engine` | string | `search_std` | Search engine type | +| `max_results` | int | 5 | Maximum number of results | ### Additional Web Settings @@ -134,6 +134,28 @@ Baidu Search uses the [Qianfan AI Search API](https://cloud.baidu.com/doc/qianfa | `prefer_native` | bool | true | Prefer provider's native search over configured search engines | | `private_host_whitelist` | string[] | `[]` | Private/internal hosts allowed for web fetching | +### `web_search` Tool Parameters + +At runtime, the `web_search` tool accepts the following parameters: + +| Field | Type | Required | Description | +|-------|------|----------|-------------| +| `query` | string | yes | Search query string | +| `count` | integer | no | Number of results to return. Default: `10`, max: `10` | +| `range` | string | no | Optional time filter: `d` (day), `w` (week), `m` (month), `y` (year) | + +If `range` is omitted, PicoClaw performs an unrestricted search. + +### Example `web_search` Call + +```json +{ + "query": "ai agent news", + "count": 10, + "range": "w" +} +``` + ## Exec Tool The exec tool is used to execute shell commands. diff --git a/docs/zh/chat-apps.md b/docs/zh/chat-apps.md index 4d1451d68..47add38ac 100644 --- a/docs/zh/chat-apps.md +++ b/docs/zh/chat-apps.md @@ -6,7 +6,7 @@ PicoClaw 支持多种聊天平台，使您的 Agent 能够连接到任何地方。 -> **注意**: 所有 Webhook 类渠道（LINE、WeCom 等）均挂载在同一个 Gateway HTTP 服务器上（`gateway.host`:`gateway.port`，默认 `127.0.0.1:18790`），无需为每个渠道单独配置端口。注意：飞书（Feishu）使用 WebSocket/SDK 模式，不通过该共享 HTTP webhook 服务器接收消息。 +> **注意**: 依赖 HTTP 回调的渠道共用同一个 Gateway HTTP 服务器（`gateway.host`:`gateway.port`，默认 `127.0.0.1:18790`），无需为每个渠道单独配置端口。飞书、钉钉、企业微信这类 Socket/Stream 模式渠道不依赖共享 webhook 服务器来接收入站消息。 ### 核心渠道 @@ -21,7 +21,7 @@ PicoClaw 支持多种聊天平台，使您的 Agent 能够连接到任何地方 | **QQ** | ⭐⭐ 中等 | 官方机器人 API，适合国内社群 | [查看文档](../channels/qq/README.zh.md) | | **钉钉 (DingTalk)** | ⭐⭐ 中等 | Stream 模式无需公网，企业办公首选 | [查看文档](../channels/dingtalk/README.zh.md) | | **LINE** | ⭐⭐⭐ 较难 | 需要 HTTPS Webhook | [查看文档](../channels/line/README.zh.md) | -| **企业微信 (WeCom)** | ⭐⭐⭐ 较难 | 支持群机器人(Webhook)、自建应用(API)和智能机器人(AI Bot) | [Bot 文档](../channels/wecom/wecom_bot/README.zh.md) / [App 文档](../channels/wecom/wecom_app/README.zh.md) / [AI Bot 文档](../channels/wecom/wecom_aibot/README.zh.md) | +| **企业微信 (WeCom)** | ⭐⭐⭐ 较难 | 官方 AI Bot WebSocket 接入，支持流式回复和媒体消息 | [查看文档](../channels/wecom/README.zh.md) | | **飞书 (Feishu)** | ⭐⭐⭐ 较难 | 企业级协作，功能丰富 | [查看文档](../channels/feishu/README.zh.md) | | **IRC** | ⭐⭐ 中等 | 服务器 + TLS 配置 | [查看文档](#irc) | | **OneBot** | ⭐⭐ 中等 | 兼容 NapCat/Go-CQHTTP，社区生态丰富 | [查看文档](../channels/onebot/README.zh.md) | @@ -492,102 +492,34 @@ picoclaw gateway
企业微信 (WeCom) -PicoClaw 支持三种企业微信集成方式： +PicoClaw 现在将企业微信统一为一个基于 WebSocket 的 AI Bot 渠道。 +它不再需要公网 webhook 回调地址。 -**方式 1: 群机器人 (Bot)** — 设置简单，支持群聊 -**方式 2: 自建应用 (App)** — 功能更多，支持主动推送，仅私聊 -**方式 3: 智能机器人 (AI Bot)** — 官方 AI Bot，流式回复，支持群聊和私聊 +完整配置说明和迁移说明请参考 [企业微信配置指南](../channels/wecom/README.zh.md)。 -详细设置请参考 [企业微信 AI Bot 配置指南](../channels/wecom/wecom_aibot/README.zh.md)。 +**推荐快速接入** -**快速设置 — 群机器人：** +**1. 认证** -**1. 创建 Bot** +```bash +picoclaw auth wecom +``` -* 企业微信管理后台 → 群聊 → 添加群机器人 -* 复制 Webhook URL（格式：`https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=xxx`） +该命令会显示二维码，等待你在企业微信里确认，然后把 `bot_id` 和 `secret` 写入 `channels.wecom`。 -**2. 配置** +**2. 如需手动配置** ```json { "channels": { "wecom": { "enabled": true, - "token": "YOUR_TOKEN", - "encoding_aes_key": "YOUR_ENCODING_AES_KEY", - "webhook_url": "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=YOUR_KEY", - "webhook_path": "/webhook/wecom", - "allow_from": [] - } - } -} -``` - -> WeCom Webhook 挂载在共享 Gateway 服务器上（`gateway.host`:`gateway.port`，默认 `127.0.0.1:18790`）。 - -**快速设置 — 自建应用：** - -**1. 创建应用** - -* 企业微信管理后台 → 应用管理 → 创建应用 -* 复制 **AgentId** 和 **Secret** -* 前往"我的企业"页面，复制 **CorpID** - -**2. 配置接收消息** - -* 在应用详情中，点击"接收消息" → "设置 API" -* 设置 URL 为 `http://your-server:18790/webhook/wecom-app` -* 生成 **Token** 和 **EncodingAESKey** - -**3. 配置** - -```json -{ - "channels": { - "wecom_app": { - "enabled": true, - "corp_id": "wwxxxxxxxxxxxxxxxx", - "corp_secret": "YOUR_CORP_SECRET", - "agent_id": 1000002, - "token": "YOUR_TOKEN", - "encoding_aes_key": "YOUR_ENCODING_AES_KEY", - "webhook_path": "/webhook/wecom-app", - "allow_from": [] - } - } -} -``` - -**4. 运行** - -```bash -picoclaw gateway -``` - -> **注意**: WeCom Webhook 回调挂载在 Gateway 端口（默认 18790）。使用反向代理配置 HTTPS。 - -**快速设置 — 智能机器人 (AI Bot)：** - -**1. 创建 AI Bot** - -* 企业微信管理后台 → 应用管理 → AI Bot -* 在 AI Bot 设置中配置回调 URL：`http://your-server:18790/webhook/wecom-aibot` -* 复制 **Token** 并点击"随机生成" **EncodingAESKey** - -**2. 配置** - -```json -{ - "channels": { - "wecom_aibot": { - "enabled": true, - "token": "YOUR_TOKEN", - "encoding_aes_key": "YOUR_43_CHAR_ENCODING_AES_KEY", - "webhook_path": "/webhook/wecom-aibot", + "bot_id": "YOUR_BOT_ID", + "secret": "YOUR_SECRET", + "websocket_url": "wss://openws.work.weixin.qq.com", + "send_thinking_message": true, "allow_from": [], - "welcome_message": "你好！有什么可以帮你的？", - "processing_message": "⏳ Processing, please wait. The results will be sent shortly." + "reasoning_channel_id": "" } } } @@ -599,7 +531,7 @@ picoclaw gateway picoclaw gateway ``` -> **注意**: 企业微信 AI Bot 使用流式拉取协议，无回复超时问题。长任务（>30 秒）会自动切换到 `response_url` 推送投递。 +> 这个分支中旧的 `wecom_app` 和 `wecom_aibot` 配置已经被统一的 `channels.wecom` 替代。
diff --git a/docs/zh/providers.md b/docs/zh/providers.md index e7b323ebf..4bcf7087e 100644 --- a/docs/zh/providers.md +++ b/docs/zh/providers.md @@ -26,6 +26,7 @@ | `mistral` | LLM (Mistral 直连) | [console.mistral.ai](https://console.mistral.ai) | | `longcat` | LLM (Longcat 直连) | [longcat.ai](https://longcat.ai) | | `modelscope` | LLM (ModelScope 直连) | [modelscope.cn](https://modelscope.cn) | +| `mimo` | LLM (小米 MiMo 直连) | [platform.xiaomimimo.com](https://platform.xiaomimimo.com) | ### 模型配置 (model_list) @@ -62,6 +63,7 @@ | **Vivgrid** | `vivgrid/` | `https://api.vivgrid.com/v1` | OpenAI | [获取密钥](https://vivgrid.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) | +| **小米 MiMo** | `mimo/` | `https://api.xiaomimimo.com/v1` | OpenAI | [获取密钥](https://platform.xiaomimimo.com) | | **Antigravity** | `antigravity/` | Google Cloud | 自定义 | 仅 OAuth | | **GitHub Copilot** | `github-copilot/` | `localhost:4321` | gRPC | - | @@ -257,6 +259,45 @@ PicoClaw 在发送请求前仅去除外层 `litellm/` 前缀，因此 `litellm/l } ``` +#### 自动模型失败切换（Cascade） + +当你在 Agent 的模型设置里配置 `primary` + `fallbacks` 时，PicoClaw 已经支持自动失败切换。 +运行时 fallback 链会在可重试错误时切到下一个候选（例如 HTTP `429`、配额/限流错误、超时错误）。 +同时会对每个候选应用 cooldown，避免对刚失败的目标立即重试。 + +```json +{ + "model_list": [ + { + "model_name": "qwen-main", + "model": "openai/qwen3.5:cloud", + "api_base": "https://api.example.com/v1", + "api_key": "sk-main" + }, + { + "model_name": "deepseek-backup", + "model": "deepseek/deepseek-chat", + "api_key": "sk-backup-1" + }, + { + "model_name": "gemini-backup", + "model": "gemini/gemini-2.5-flash", + "api_key": "sk-backup-2" + } + ], + "agents": { + "defaults": { + "model": { + "primary": "qwen-main", + "fallbacks": ["deepseek-backup", "gemini-backup"] + } + } + } +} +``` + +如果你在同一模型上启用了 key 级失败切换，PicoClaw 会先在该模型的多 key 候选间切换，再继续切到跨模型备选。 + #### 从旧的 `providers` 配置迁移 旧的 `providers` 配置格式**已弃用**，但为向后兼容仍支持。 diff --git a/go.mod b/go.mod index bce41d0d3..54c275102 100644 --- a/go.mod +++ b/go.mod @@ -12,6 +12,7 @@ require ( github.com/aws/aws-sdk-go-v2/service/bedrockruntime v1.50.2 github.com/bwmarrin/discordgo v0.29.0 github.com/caarlos0/env/v11 v11.4.0 + github.com/creack/pty v1.1.24 github.com/ergochat/irc-go v0.6.0 github.com/ergochat/readline v0.1.3 github.com/gdamore/tcell/v2 v2.13.8 @@ -31,6 +32,7 @@ require ( github.com/spf13/cobra v1.10.2 github.com/stretchr/testify v1.11.1 github.com/tencent-connect/botgo v0.2.1 + go.mau.fi/util v0.9.7 go.mau.fi/whatsmeow v0.0.0-20260219150138-7ae702b1eed4 golang.org/x/oauth2 v0.36.0 golang.org/x/term v0.41.0 @@ -39,6 +41,7 @@ require ( gopkg.in/yaml.v3 v3.0.1 maunium.net/go/mautrix v0.26.4 modernc.org/sqlite v1.46.1 + rsc.io/qr v0.2.0 ) require ( @@ -67,6 +70,7 @@ require ( github.com/lucasb-eyer/go-colorful v1.3.0 // indirect github.com/mattn/go-colorable v0.1.14 // indirect github.com/mattn/go-isatty v0.0.20 // indirect + github.com/mattn/go-sqlite3 v1.14.34 // indirect github.com/ncruces/go-strftime v1.0.0 // indirect github.com/petermattis/goid v0.0.0-20260226131333-17d1149c6ac6 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect @@ -77,13 +81,11 @@ require ( github.com/spf13/pflag v1.0.10 // indirect github.com/vektah/gqlparser/v2 v2.5.27 // indirect go.mau.fi/libsignal v0.2.1 // indirect - go.mau.fi/util v0.9.7 // indirect golang.org/x/exp v0.0.0-20260312153236-7ab1446f8b90 // indirect golang.org/x/text v0.35.0 // indirect modernc.org/libc v1.67.6 // indirect modernc.org/mathutil v1.7.1 // indirect modernc.org/memory v1.11.0 // indirect - rsc.io/qr v0.2.0 // indirect ) require ( diff --git a/go.sum b/go.sum index 87117bc98..ae12473f3 100644 --- a/go.sum +++ b/go.sum @@ -70,6 +70,8 @@ github.com/coder/websocket v1.8.14/go.mod h1:NX3SzP+inril6yawo5CQXx8+fk145lPDC6p github.com/coreos/go-systemd/v22 v22.5.0/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc= github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g= github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= +github.com/creack/pty v1.1.24 h1:bJrF4RRfyJnbTJqzRLHzcGaZK1NeM5kTC9jGgovnR1s= +github.com/creack/pty v1.1.24/go.mod h1:08sCNb52WyoAwi2QDyzUCTgcvVFhUzewun7wtTfvcwE= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= diff --git a/pkg/agent/context.go b/pkg/agent/context.go index 12e3cdd4d..c3fcc9fff 100644 --- a/pkg/agent/context.go +++ b/pkg/agent/context.go @@ -26,6 +26,7 @@ type ContextBuilder struct { memory *MemoryStore toolDiscoveryBM25 bool toolDiscoveryRegex bool + splitOnMarker bool // Cache for system prompt to avoid rebuilding on every call. // This fixes issue #607: repeated reprocessing of the entire context. @@ -52,6 +53,11 @@ func (cb *ContextBuilder) WithToolDiscovery(useBM25, useRegex bool) *ContextBuil return cb } +func (cb *ContextBuilder) WithSplitOnMarker(enabled bool) *ContextBuilder { + cb.splitOnMarker = enabled + return cb +} + func getGlobalConfigDir() string { if home := os.Getenv(config.EnvHome); home != "" { return home @@ -157,6 +163,14 @@ The following skills extend your capabilities. To use a skill, read its SKILL.md parts = append(parts, "# Memory\n\n"+memoryContext) } + // Multi-Message Sending (if enabled) + if cb.splitOnMarker { + parts = append(parts, `# MULTI-MESSAGE OUTPUT +You MUST frequently use <|[SPLIT]|> to break your responses into multiple short messages. NEVER output a single long wall of text. Actively split distinct concepts or parts. Example: Message part 1<|[SPLIT]|>Message part 2<|[SPLIT]|>Message part 3 + +Each part separated by the marker will be sent as an independent message.`) + } + // Join with "---" separator return strings.Join(parts, "\n\n---\n\n") } diff --git a/pkg/agent/instance.go b/pkg/agent/instance.go index 34d401186..cef736981 100644 --- a/pkg/agent/instance.go +++ b/pkg/agent/instance.go @@ -103,10 +103,12 @@ func NewAgentInstance( sessions := initSessionStore(sessionsDir) mcpDiscoveryActive := cfg.Tools.MCP.Enabled && cfg.Tools.MCP.Discovery.Enabled - contextBuilder := NewContextBuilder(workspace).WithToolDiscovery( - mcpDiscoveryActive && cfg.Tools.MCP.Discovery.UseBM25, - mcpDiscoveryActive && cfg.Tools.MCP.Discovery.UseRegex, - ) + contextBuilder := NewContextBuilder(workspace). + WithToolDiscovery( + mcpDiscoveryActive && cfg.Tools.MCP.Discovery.UseBM25, + mcpDiscoveryActive && cfg.Tools.MCP.Discovery.UseRegex, + ). + WithSplitOnMarker(cfg.Agents.Defaults.SplitOnMarker) agentID := routing.DefaultAgentID agentName := "" diff --git a/pkg/agent/instance_test.go b/pkg/agent/instance_test.go index e073cb929..e296a18cb 100644 --- a/pkg/agent/instance_test.go +++ b/pkg/agent/instance_test.go @@ -236,8 +236,9 @@ func TestNewAgentInstance_AllowsMediaTempDirForReadListAndExec(t *testing.T) { t.Fatal("exec tool not registered") } execResult := execTool.Execute(context.Background(), map[string]any{ - "command": "cat " + filepath.Base(mediaPath), - "working_dir": mediaDir, + "action": "run", + "command": "cat " + filepath.Base(mediaPath), + "cwd": mediaDir, }) if execResult.IsError { t.Fatalf("exec should allow media temp dir, got: %s", execResult.ForLLM) diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 725d42614..81b979490 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -85,6 +85,7 @@ type processOptions struct { DefaultResponse string // Response when LLM returns empty EnableSummary bool // Whether to trigger summarization SendResponse bool // Whether to send response via bus + SuppressToolFeedback bool // Whether to suppress inline tool feedback messages NoHistory bool // If true, don't load session history (for heartbeat) SkipInitialSteeringPoll bool // If true, skip the steering poll at loop start (used by Continue) } @@ -96,14 +97,15 @@ type continuationTarget struct { } const ( - defaultResponse = "The model returned an empty response. This may indicate a provider error or token limit." - toolLimitResponse = "I've reached `max_tool_iterations` without a final response. Increase `max_tool_iterations` in config.json if this task needs more tool steps." - sessionKeyAgentPrefix = "agent:" - metadataKeyAccountID = "account_id" - metadataKeyGuildID = "guild_id" - metadataKeyTeamID = "team_id" - metadataKeyParentPeerKind = "parent_peer_kind" - metadataKeyParentPeerID = "parent_peer_id" + defaultResponse = "The model returned an empty response. This may indicate a provider error or token limit." + toolLimitResponse = "I've reached `max_tool_iterations` without a final response. Increase `max_tool_iterations` in config.json if this task needs more tool steps." + handledToolResponseSummary = "Requested output delivered via tool attachment." + sessionKeyAgentPrefix = "agent:" + metadataKeyAccountID = "account_id" + metadataKeyGuildID = "guild_id" + metadataKeyTeamID = "team_id" + metadataKeyParentPeerKind = "parent_peer_kind" + metadataKeyParentPeerID = "parent_peer_id" ) func NewAgentLoop( @@ -1030,13 +1032,13 @@ func (al *AgentLoop) GetConfig() *config.Config { func (al *AgentLoop) SetMediaStore(s media.MediaStore) { al.mediaStore = s - // Propagate store to send_file tools in all agents. + // Propagate store to all registered tools that can emit media. registry := al.GetRegistry() - registry.ForEachTool("send_file", func(t tools.Tool) { - if sf, ok := t.(*tools.SendFileTool); ok { - sf.SetMediaStore(s) + for _, agentID := range registry.ListAgentIDs() { + if agent, ok := registry.GetAgent(agentID); ok { + agent.Tools.SetMediaStore(s) } - }) + } } // SetTranscriber injects a voice transcriber for agent-level audio transcription. @@ -1241,14 +1243,15 @@ func (al *AgentLoop) ProcessHeartbeat( return "", fmt.Errorf("no default agent for heartbeat") } return al.runAgentLoop(ctx, agent, processOptions{ - SessionKey: "heartbeat", - Channel: channel, - ChatID: chatID, - UserMessage: content, - DefaultResponse: defaultResponse, - EnableSummary: false, - SendResponse: false, - NoHistory: true, // Don't load session history for heartbeat + SessionKey: "heartbeat", + Channel: channel, + ChatID: chatID, + UserMessage: content, + DefaultResponse: defaultResponse, + EnableSummary: false, + SendResponse: false, + SuppressToolFeedback: true, + NoHistory: true, // Don't load session history for heartbeat }) } @@ -1946,6 +1949,7 @@ turnLoop: isContextError := !isTimeoutError && (strings.Contains(errMsg, "context_length_exceeded") || strings.Contains(errMsg, "context window") || + strings.Contains(errMsg, "context_window") || strings.Contains(errMsg, "maximum context length") || strings.Contains(errMsg, "token limit") || strings.Contains(errMsg, "too many tokens") || @@ -2165,6 +2169,7 @@ turnLoop: "iteration": iteration, }) + allResponsesHandled := len(normalizedToolCalls) > 0 assistantMsg := providers.Message{ Role: "assistant", Content: response.Content, @@ -2221,6 +2226,7 @@ turnLoop: toolArgs = toolReq.Arguments } case HookActionDenyTool: + allResponsesHandled = false denyContent := hookDeniedToolContent("Tool execution denied by hook", decision.Reason) al.emitEvent( EventKindToolExecSkipped, @@ -2260,6 +2266,7 @@ turnLoop: ChatID: ts.chatID, }) if !approval.Approved { + allResponsesHandled = false denyContent := hookDeniedToolContent("Tool execution denied by approval hook", approval.Reason) al.emitEvent( EventKindToolExecSkipped, @@ -2301,7 +2308,9 @@ turnLoop: ) // Send tool feedback to chat channel if enabled (from HEAD) - if al.cfg.Agents.Defaults.IsToolFeedbackEnabled() && ts.channel != "" { + if al.cfg.Agents.Defaults.IsToolFeedbackEnabled() && + ts.channel != "" && + !ts.opts.SuppressToolFeedback { feedbackPreview := utils.Truncate( string(argsJSON), al.cfg.Agents.Defaults.GetToolFeedbackMaxArgsLength(), @@ -2333,10 +2342,7 @@ turnLoop: } // Determine content for the agent loop (ForLLM or error). - content := result.ForLLM - if content == "" && result.Err != nil { - content = result.Err.Error() - } + content := result.ContentForLLM() if content == "" { return } @@ -2420,6 +2426,50 @@ turnLoop: if toolResult == nil { toolResult = tools.ErrorResult("hook returned nil tool result") } + if len(toolResult.Media) > 0 && toolResult.ResponseHandled { + parts := make([]bus.MediaPart, 0, len(toolResult.Media)) + for _, ref := range toolResult.Media { + part := bus.MediaPart{Ref: ref} + if al.mediaStore != nil { + if _, meta, err := al.mediaStore.ResolveWithMeta(ref); err == nil { + part.Filename = meta.Filename + part.ContentType = meta.ContentType + part.Type = inferMediaType(meta.Filename, meta.ContentType) + } + } + parts = append(parts, part) + } + outboundMedia := bus.OutboundMediaMessage{ + Channel: ts.channel, + ChatID: ts.chatID, + Parts: parts, + } + if al.channelManager != nil && ts.channel != "" && !constants.IsInternalChannel(ts.channel) { + if err := al.channelManager.SendMedia(ctx, outboundMedia); err != nil { + logger.WarnCF("agent", "Failed to deliver handled tool media", + map[string]any{ + "agent_id": ts.agent.ID, + "tool": toolName, + "channel": ts.channel, + "chat_id": ts.chatID, + "error": err.Error(), + }) + toolResult = tools.ErrorResult(fmt.Sprintf("failed to deliver attachment: %v", err)).WithError(err) + } + } else if al.bus != nil { + al.bus.PublishOutboundMedia(ctx, outboundMedia) + // Queuing media is only best-effort; it has not been delivered yet. + toolResult.ResponseHandled = false + } + } + + if len(toolResult.Media) > 0 && !toolResult.ResponseHandled { + toolResult.ArtifactTags = buildArtifactTags(al.mediaStore, toolResult.Media) + } + + if !toolResult.ResponseHandled { + allResponsesHandled = false + } if !toolResult.Silent && toolResult.ForUser != "" && ts.opts.SendResponse { al.bus.PublishOutbound(ctx, bus.OutboundMessage{ @@ -2434,30 +2484,7 @@ turnLoop: }) } - if len(toolResult.Media) > 0 { - parts := make([]bus.MediaPart, 0, len(toolResult.Media)) - for _, ref := range toolResult.Media { - part := bus.MediaPart{Ref: ref} - if al.mediaStore != nil { - if _, meta, err := al.mediaStore.ResolveWithMeta(ref); err == nil { - part.Filename = meta.Filename - part.ContentType = meta.ContentType - part.Type = inferMediaType(meta.Filename, meta.ContentType) - } - } - parts = append(parts, part) - } - al.bus.PublishOutboundMedia(ctx, bus.OutboundMediaMessage{ - Channel: ts.channel, - ChatID: ts.chatID, - Parts: parts, - }) - } - - contentForLLM := toolResult.ForLLM - if contentForLLM == "" && toolResult.Err != nil { - contentForLLM = toolResult.Err.Error() - } + contentForLLM := toolResult.ContentForLLM() // Filter sensitive data (API keys, tokens, secrets) before sending to LLM if al.cfg.Tools.IsFilterSensitiveDataEnabled() { @@ -2552,6 +2579,70 @@ turnLoop: } } + if allResponsesHandled { + if len(pendingMessages) > 0 { + logger.InfoCF("agent", "Pending steering exists after handled tool delivery; continuing turn before finalizing", + map[string]any{ + "agent_id": ts.agent.ID, + "steering_count": len(pendingMessages), + "session_key": ts.sessionKey, + }) + finalContent = "" + goto turnLoop + } + + if steerMsgs := al.dequeueSteeringMessagesForScope(ts.sessionKey); len(steerMsgs) > 0 { + logger.InfoCF("agent", "Steering arrived after handled tool delivery; continuing turn before finalizing", + map[string]any{ + "agent_id": ts.agent.ID, + "steering_count": len(steerMsgs), + "session_key": ts.sessionKey, + }) + pendingMessages = append(pendingMessages, steerMsgs...) + finalContent = "" + goto turnLoop + } + + summaryMsg := providers.Message{ + Role: "assistant", + Content: handledToolResponseSummary, + } + + if !ts.opts.NoHistory { + ts.agent.Sessions.AddMessage(ts.sessionKey, summaryMsg.Role, summaryMsg.Content) + ts.recordPersistedMessage(summaryMsg) + if err := ts.agent.Sessions.Save(ts.sessionKey); err != nil { + turnStatus = TurnEndStatusError + al.emitEvent( + EventKindError, + ts.eventMeta("runTurn", "turn.error"), + ErrorPayload{ + Stage: "session_save", + Message: err.Error(), + }, + ) + return turnResult{}, err + } + } + if ts.opts.EnableSummary { + al.maybeSummarize(ts.agent, ts.sessionKey, ts.scope) + } + + ts.setPhase(TurnPhaseCompleted) + ts.setFinalContent("") + logger.InfoCF("agent", "Tool output satisfied delivery; ending turn without follow-up LLM", + map[string]any{ + "agent_id": ts.agent.ID, + "iteration": iteration, + "tool_count": len(normalizedToolCalls), + }) + return turnResult{ + finalContent: "", + status: turnStatus, + followUps: append([]bus.InboundMessage(nil), ts.followUps...), + }, nil + } + ts.agent.Tools.TickTTL() logger.DebugCF("agent", "TTL tick after tool execution", map[string]any{ "agent_id": ts.agent.ID, "iteration": iteration, @@ -3159,6 +3250,97 @@ func (al *AgentLoop) handleCommand( } } +func activeSkillNames(agent *AgentInstance, opts processOptions) []string { + if agent == nil { + return nil + } + + combined := make([]string, 0, len(agent.SkillsFilter)+len(opts.ForcedSkills)) + combined = append(combined, agent.SkillsFilter...) + combined = append(combined, opts.ForcedSkills...) + if len(combined) == 0 { + return nil + } + + var resolved []string + seen := make(map[string]struct{}, len(combined)) + for _, name := range combined { + name = strings.TrimSpace(name) + if name == "" { + continue + } + if agent.ContextBuilder != nil { + if canonical, ok := agent.ContextBuilder.ResolveSkillName(name); ok { + name = canonical + } + } + key := strings.ToLower(name) + if _, ok := seen[key]; ok { + continue + } + seen[key] = struct{}{} + resolved = append(resolved, name) + } + + return resolved +} + +func (al *AgentLoop) applyExplicitSkillCommand( + raw string, + agent *AgentInstance, + opts *processOptions, +) (matched bool, handled bool, reply string) { + cmdName, ok := commands.CommandName(raw) + if !ok || cmdName != "use" { + return false, false, "" + } + + if agent == nil || agent.ContextBuilder == nil { + return true, true, commandsUnavailableSkillMessage() + } + + parts := strings.Fields(strings.TrimSpace(raw)) + if len(parts) < 2 { + return true, true, buildUseCommandHelp(agent) + } + + arg := strings.TrimSpace(parts[1]) + if strings.EqualFold(arg, "clear") || strings.EqualFold(arg, "off") { + if opts != nil { + al.clearPendingSkills(opts.SessionKey) + } + return true, true, "Cleared pending skill override." + } + + skillName, ok := agent.ContextBuilder.ResolveSkillName(arg) + if !ok { + return true, true, fmt.Sprintf("Unknown skill: %s\nUse /list skills to see installed skills.", arg) + } + + if len(parts) < 3 { + if opts == nil || strings.TrimSpace(opts.SessionKey) == "" { + return true, true, commandsUnavailableSkillMessage() + } + al.setPendingSkills(opts.SessionKey, []string{skillName}) + return true, true, fmt.Sprintf( + "Skill %q is armed for your next message. Send your next prompt normally, or use /use clear to cancel.", + skillName, + ) + } + + message := strings.TrimSpace(strings.Join(parts[2:], " ")) + if message == "" { + return true, true, buildUseCommandHelp(agent) + } + + if opts != nil { + opts.ForcedSkills = append(opts.ForcedSkills, skillName) + opts.UserMessage = message + } + + return true, false, "" +} + func (al *AgentLoop) buildCommandsRuntime(agent *AgentInstance, opts *processOptions) *commands.Runtime { registry := al.GetRegistry() cfg := al.GetConfig() @@ -3199,6 +3381,9 @@ func (al *AgentLoop) buildCommandsRuntime(agent *AgentInstance, opts *processOpt return al.reloadFunc() } if agent != nil { + if agent.ContextBuilder != nil { + rt.ListSkillNames = agent.ContextBuilder.ListSkillNames + } rt.GetModelInfo = func() (string, string) { return agent.Model, resolvedCandidateProvider(agent.Candidates, cfg.Agents.Defaults.Provider) } @@ -3251,79 +3436,6 @@ func (al *AgentLoop) buildCommandsRuntime(agent *AgentInstance, opts *processOpt return rt } -func activeSkillNames(agent *AgentInstance, opts processOptions) []string { - var out []string - seen := make(map[string]struct{}) - - appendNames := func(names []string) { - for _, name := range names { - name = strings.TrimSpace(name) - if name == "" { - continue - } - if _, exists := seen[name]; exists { - continue - } - seen[name] = struct{}{} - out = append(out, name) - } - } - - if agent != nil { - appendNames(agent.SkillsFilter) - } - appendNames(opts.ForcedSkills) - - return out -} - -func (al *AgentLoop) applyExplicitSkillCommand( - raw string, - agent *AgentInstance, - opts *processOptions, -) (matched bool, handled bool, reply string) { - commandName, ok := commands.CommandName(raw) - if !ok || commandName != "use" { - return false, false, "" - } - - if agent == nil || agent.ContextBuilder == nil { - return true, true, commandsUnavailableSkillMessage() - } - - fields := strings.Fields(strings.TrimSpace(raw)) - if len(fields) < 2 { - return true, true, buildUseCommandHelp(agent) - } - - if strings.EqualFold(fields[1], "clear") || strings.EqualFold(fields[1], "off") { - al.clearPendingSkills(opts.SessionKey) - return true, true, "Cleared pending skill override." - } - - canonicalSkill, ok := agent.ContextBuilder.ResolveSkillName(fields[1]) - if !ok { - return true, true, fmt.Sprintf("Unknown skill: %s\nUse /list skills to see installed skills.", fields[1]) - } - - if len(fields) == 2 { - al.setPendingSkills(opts.SessionKey, []string{canonicalSkill}) - return true, true, fmt.Sprintf( - "Skill %q is armed for your next message.\nSend your next request normally, or use /use clear to cancel.", - canonicalSkill, - ) - } - - message := strings.TrimSpace(strings.Join(fields[2:], " ")) - if message == "" { - return true, true, buildUseCommandHelp(agent) - } - - opts.UserMessage = message - opts.ForcedSkills = append(opts.ForcedSkills, canonicalSkill) - return true, false, "" -} - func commandsUnavailableSkillMessage() string { return "Skill selection is unavailable in the current context." } diff --git a/pkg/agent/loop_media.go b/pkg/agent/loop_media.go index 1380f0214..e8314c10d 100644 --- a/pkg/agent/loop_media.go +++ b/pkg/agent/loop_media.go @@ -87,6 +87,24 @@ func resolveMediaRefs(messages []providers.Message, store media.MediaStore, maxS return result } +func buildArtifactTags(store media.MediaStore, refs []string) []string { + if store == nil || len(refs) == 0 { + return nil + } + + tags := make([]string, 0, len(refs)) + for _, ref := range refs { + localPath, meta, err := store.ResolveWithMeta(ref) + if err != nil { + continue + } + mime := detectMIME(localPath, meta) + tags = append(tags, buildPathTag(mime, localPath)) + } + + return tags +} + // detectMIME determines the MIME type from metadata or magic-bytes detection. // Returns empty string if detection fails. func detectMIME(localPath string, meta media.MediaMeta) string { diff --git a/pkg/agent/loop_test.go b/pkg/agent/loop_test.go index 1a4a44edf..2366b1277 100644 --- a/pkg/agent/loop_test.go +++ b/pkg/agent/loop_test.go @@ -3,6 +3,7 @@ package agent import ( "context" "encoding/json" + "errors" "fmt" "net/http" "net/http/httptest" @@ -33,6 +34,41 @@ func (f *fakeChannel) IsAllowed(string) bool { func (f *fakeChannel) IsAllowedSender(sender bus.SenderInfo) bool { return true } func (f *fakeChannel) ReasoningChannelID() string { return f.id } +type fakeMediaChannel struct { + fakeChannel + sentMedia []bus.OutboundMediaMessage +} + +func (f *fakeMediaChannel) SendMedia(ctx context.Context, msg bus.OutboundMediaMessage) error { + f.sentMedia = append(f.sentMedia, msg) + return nil +} + +func newStartedTestChannelManager( + t *testing.T, + msgBus *bus.MessageBus, + store media.MediaStore, + name string, + ch channels.Channel, +) *channels.Manager { + t.Helper() + + cm, err := channels.NewManager(&config.Config{}, msgBus, store) + if err != nil { + t.Fatalf("NewManager() error = %v", err) + } + cm.RegisterChannel(name, ch) + if err := cm.StartAll(context.Background()); err != nil { + t.Fatalf("StartAll() error = %v", err) + } + t.Cleanup(func() { + if err := cm.StopAll(context.Background()); err != nil { + t.Fatalf("StopAll() error = %v", err) + } + }) + return cm +} + type recordingProvider struct { lastMessages []providers.Message } @@ -289,6 +325,86 @@ func TestProcessMessage_UseCommandArmsSkillForNextMessage(t *testing.T) { } } +func TestApplyExplicitSkillCommand_ArmsSkillForNextMessage(t *testing.T) { + al, cfg, _, _, cleanup := newTestAgentLoop(t) + defer cleanup() + + if err := os.MkdirAll(filepath.Join(cfg.Agents.Defaults.Workspace, "skills", "finance-news"), 0o755); err != nil { + t.Fatalf("MkdirAll(skill) error = %v", err) + } + if err := os.WriteFile( + filepath.Join(cfg.Agents.Defaults.Workspace, "skills", "finance-news", "SKILL.md"), + []byte("# Finance News\n\nUse web tools for current finance updates.\n"), + 0o644, + ); err != nil { + t.Fatalf("WriteFile(SKILL.md) error = %v", err) + } + + agent := al.GetRegistry().GetDefaultAgent() + if agent == nil { + t.Fatal("expected default agent") + } + + opts := &processOptions{SessionKey: "agent:main:test"} + matched, handled, reply := al.applyExplicitSkillCommand("/use finance-news", agent, opts) + if !matched { + t.Fatal("expected /use command to match") + } + if !handled { + t.Fatal("expected /use without inline message to be handled immediately") + } + if !strings.Contains(reply, `Skill "finance-news" is armed for your next message`) { + t.Fatalf("unexpected reply: %q", reply) + } + + pending := al.takePendingSkills(opts.SessionKey) + if len(pending) != 1 || pending[0] != "finance-news" { + t.Fatalf("pending skills = %#v, want [finance-news]", pending) + } +} + +func TestApplyExplicitSkillCommand_InlineMessageMutatesOptions(t *testing.T) { + al, cfg, _, _, cleanup := newTestAgentLoop(t) + defer cleanup() + + if err := os.MkdirAll(filepath.Join(cfg.Agents.Defaults.Workspace, "skills", "finance-news"), 0o755); err != nil { + t.Fatalf("MkdirAll(skill) error = %v", err) + } + if err := os.WriteFile( + filepath.Join(cfg.Agents.Defaults.Workspace, "skills", "finance-news", "SKILL.md"), + []byte("# Finance News\n\nUse web tools for current finance updates.\n"), + 0o644, + ); err != nil { + t.Fatalf("WriteFile(SKILL.md) error = %v", err) + } + + agent := al.GetRegistry().GetDefaultAgent() + if agent == nil { + t.Fatal("expected default agent") + } + + opts := &processOptions{ + SessionKey: "agent:main:test", + UserMessage: "/use finance-news dammi le ultime news", + } + matched, handled, reply := al.applyExplicitSkillCommand(opts.UserMessage, agent, opts) + if !matched { + t.Fatal("expected /use command to match") + } + if handled { + t.Fatal("expected /use with inline message to fall through into normal agent execution") + } + if reply != "" { + t.Fatalf("unexpected reply: %q", reply) + } + if opts.UserMessage != "dammi le ultime news" { + t.Fatalf("opts.UserMessage = %q, want %q", opts.UserMessage, "dammi le ultime news") + } + if len(opts.ForcedSkills) != 1 || opts.ForcedSkills[0] != "finance-news" { + t.Fatalf("opts.ForcedSkills = %#v, want [finance-news]", opts.ForcedSkills) + } +} + func TestRecordLastChannel(t *testing.T) { al, cfg, msgBus, provider, cleanup := newTestAgentLoop(t) defer cleanup() @@ -455,6 +571,217 @@ func TestToolRegistry_GetDefinitions(t *testing.T) { } } +func TestProcessMessage_MediaToolHandledSkipsFollowUpLLMAndFinalText(t *testing.T) { + tmpDir := t.TempDir() + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + ModelName: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + }, + }, + } + + msgBus := bus.NewMessageBus() + provider := &handledMediaProvider{} + al := NewAgentLoop(cfg, msgBus, provider) + + store := media.NewFileMediaStore() + al.SetMediaStore(store) + telegramChannel := &fakeMediaChannel{fakeChannel: fakeChannel{id: "rid-telegram"}} + al.SetChannelManager(newStartedTestChannelManager(t, msgBus, store, "telegram", telegramChannel)) + + imagePath := filepath.Join(tmpDir, "screen.png") + if err := os.WriteFile(imagePath, []byte("fake screenshot"), 0o644); err != nil { + t.Fatalf("WriteFile(imagePath) error = %v", err) + } + + al.RegisterTool(&handledMediaTool{ + store: store, + path: imagePath, + }) + + response, err := al.processMessage(context.Background(), bus.InboundMessage{ + Channel: "telegram", + ChatID: "chat1", + SenderID: "user1", + Content: "take a screenshot of the screen and send it to me", + }) + if err != nil { + t.Fatalf("processMessage() error = %v", err) + } + if response != "" { + t.Fatalf("expected no final response when media tool already handled delivery, got %q", response) + } + if provider.calls != 1 { + t.Fatalf("expected exactly 1 LLM call, got %d", provider.calls) + } + if len(provider.toolCounts) != 1 { + t.Fatalf("expected tool counts for 1 provider call, got %d", len(provider.toolCounts)) + } + if provider.toolCounts[0] == 0 { + t.Fatal("expected tools to be available on the first LLM call") + } + + if len(telegramChannel.sentMedia) != 1 { + t.Fatalf("expected exactly 1 synchronously sent media message, got %d", len(telegramChannel.sentMedia)) + } + if telegramChannel.sentMedia[0].Channel != "telegram" || telegramChannel.sentMedia[0].ChatID != "chat1" { + t.Fatalf("unexpected sent media target: %+v", telegramChannel.sentMedia[0]) + } + if len(telegramChannel.sentMedia[0].Parts) != 1 { + t.Fatalf("expected exactly 1 sent media part, got %d", len(telegramChannel.sentMedia[0].Parts)) + } + + select { + case extra := <-msgBus.OutboundMediaChan(): + t.Fatalf("expected handled media to bypass async queue, got %+v", extra) + default: + } + + defaultAgent := al.GetRegistry().GetDefaultAgent() + if defaultAgent == nil { + t.Fatal("expected default agent") + } + route, _, err := al.resolveMessageRoute(bus.InboundMessage{ + Channel: "telegram", + ChatID: "chat1", + SenderID: "user1", + Content: "take a screenshot of the screen and send it to me", + }) + if err != nil { + t.Fatalf("resolveMessageRoute() error = %v", err) + } + sessionKey := resolveScopeKey(route, "") + history := defaultAgent.Sessions.GetHistory(sessionKey) + if len(history) == 0 { + t.Fatal("expected session history to be saved") + } + last := history[len(history)-1] + if last.Role != "assistant" || last.Content != "Requested output delivered via tool attachment." { + t.Fatalf("expected handled assistant summary in history, got %+v", last) + } +} + +func TestProcessMessage_HandledToolProcessesQueuedSteeringBeforeReturning(t *testing.T) { + tmpDir := t.TempDir() + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + ModelName: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + }, + }, + } + + msgBus := bus.NewMessageBus() + provider := &handledMediaWithSteeringProvider{} + al := NewAgentLoop(cfg, msgBus, provider) + + store := media.NewFileMediaStore() + al.SetMediaStore(store) + telegramChannel := &fakeMediaChannel{fakeChannel: fakeChannel{id: "rid-telegram"}} + al.SetChannelManager(newStartedTestChannelManager(t, msgBus, store, "telegram", telegramChannel)) + + imagePath := filepath.Join(tmpDir, "screen-steering.png") + if err := os.WriteFile(imagePath, []byte("fake screenshot"), 0o644); err != nil { + t.Fatalf("WriteFile(imagePath) error = %v", err) + } + + al.RegisterTool(&handledMediaWithSteeringTool{ + store: store, + path: imagePath, + loop: al, + }) + + response, err := al.processMessage(context.Background(), bus.InboundMessage{ + Channel: "telegram", + ChatID: "chat1", + SenderID: "user1", + Content: "take a screenshot of the screen and send it to me", + }) + if err != nil { + t.Fatalf("processMessage() error = %v", err) + } + if response != "Handled the queued steering message." { + t.Fatalf("response = %q, want queued steering response", response) + } + if provider.calls != 2 { + t.Fatalf("expected 2 LLM calls after queued steering, got %d", provider.calls) + } + if len(telegramChannel.sentMedia) != 1 { + t.Fatalf("expected exactly 1 synchronously sent media message, got %d", len(telegramChannel.sentMedia)) + } +} + +func TestProcessMessage_MediaArtifactCanBeForwardedBySendFile(t *testing.T) { + tmpDir := t.TempDir() + cfg := config.DefaultConfig() + cfg.Agents.Defaults.Workspace = tmpDir + cfg.Agents.Defaults.ModelName = "test-model" + cfg.Agents.Defaults.MaxTokens = 4096 + cfg.Agents.Defaults.MaxToolIterations = 10 + + msgBus := bus.NewMessageBus() + provider := &artifactThenSendProvider{} + al := NewAgentLoop(cfg, msgBus, provider) + + store := media.NewFileMediaStore() + al.SetMediaStore(store) + telegramChannel := &fakeMediaChannel{fakeChannel: fakeChannel{id: "rid-telegram"}} + al.SetChannelManager(newStartedTestChannelManager(t, msgBus, store, "telegram", telegramChannel)) + + mediaDir := media.TempDir() + if err := os.MkdirAll(mediaDir, 0o700); err != nil { + t.Fatalf("MkdirAll(mediaDir) error = %v", err) + } + imagePath := filepath.Join(mediaDir, "artifact-screen.png") + if err := os.WriteFile(imagePath, []byte("fake screenshot"), 0o644); err != nil { + t.Fatalf("WriteFile(imagePath) error = %v", err) + } + + al.RegisterTool(&mediaArtifactTool{ + store: store, + path: imagePath, + }) + + response, err := al.processMessage(context.Background(), bus.InboundMessage{ + Channel: "telegram", + ChatID: "chat1", + SenderID: "user1", + Content: "take a screenshot of the screen and send it to me", + }) + if err != nil { + t.Fatalf("processMessage() error = %v", err) + } + if response != "" { + t.Fatalf("expected no final response after send_file handled delivery, got %q", response) + } + if provider.calls != 2 { + t.Fatalf("expected 2 LLM calls (artifact + send_file), got %d", provider.calls) + } + + if len(telegramChannel.sentMedia) != 1 { + t.Fatalf("expected exactly 1 synchronously sent media message, got %d", len(telegramChannel.sentMedia)) + } + if telegramChannel.sentMedia[0].Channel != "telegram" || telegramChannel.sentMedia[0].ChatID != "chat1" { + t.Fatalf("unexpected sent media target: %+v", telegramChannel.sentMedia[0]) + } + if len(telegramChannel.sentMedia[0].Parts) != 1 { + t.Fatalf("expected exactly 1 sent media part, got %d", len(telegramChannel.sentMedia[0].Parts)) + } + + select { + case extra := <-msgBus.OutboundMediaChan(): + t.Fatalf("expected synchronous send_file delivery to bypass async queue, got %+v", extra) + default: + } +} + // TestAgentLoop_GetStartupInfo verifies startup info contains tools func TestAgentLoop_GetStartupInfo(t *testing.T) { tmpDir, err := os.MkdirTemp("", "agent-test-*") @@ -600,6 +927,132 @@ func (m *countingMockProvider) GetDefaultModel() string { return "counting-mock-model" } +type handledMediaProvider struct { + calls int + toolCounts []int +} + +func (m *handledMediaProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, +) (*providers.LLMResponse, error) { + m.calls++ + m.toolCounts = append(m.toolCounts, len(tools)) + if m.calls == 1 { + return &providers.LLMResponse{ + Content: "Taking the screenshot now.", + ToolCalls: []providers.ToolCall{{ + ID: "call_handled_media", + Type: "function", + Name: "handled_media_tool", + Arguments: map[string]any{}, + }}, + }, nil + } + return &providers.LLMResponse{}, nil +} + +func (m *handledMediaProvider) GetDefaultModel() string { + return "handled-media-model" +} + +type artifactThenSendProvider struct { + calls int +} + +func (m *artifactThenSendProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, +) (*providers.LLMResponse, error) { + m.calls++ + if m.calls == 1 { + return &providers.LLMResponse{ + Content: "Taking the screenshot now.", + ToolCalls: []providers.ToolCall{{ + ID: "call_artifact_media", + Type: "function", + Name: "media_artifact_tool", + Arguments: map[string]any{}, + }}, + }, nil + } + + var artifactPath string + for i := len(messages) - 1; i >= 0; i-- { + if messages[i].Role != "tool" { + continue + } + start := strings.Index(messages[i].Content, "[file:") + if start < 0 { + continue + } + rest := messages[i].Content[start+len("[file:"):] + end := strings.Index(rest, "]") + if end < 0 { + continue + } + artifactPath = rest[:end] + break + } + if artifactPath == "" { + return nil, fmt.Errorf("provider did not receive artifact path in tool result") + } + + return &providers.LLMResponse{ + Content: "", + ToolCalls: []providers.ToolCall{{ + ID: "call_send_file", + Type: "function", + Name: "send_file", + Arguments: map[string]any{"path": artifactPath}, + }}, + }, nil +} + +func (m *artifactThenSendProvider) GetDefaultModel() string { + return "artifact-then-send-model" +} + +type toolFeedbackProvider struct { + filePath string + calls int +} + +func (m *toolFeedbackProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, +) (*providers.LLMResponse, error) { + m.calls++ + if m.calls == 1 { + return &providers.LLMResponse{ + ToolCalls: []providers.ToolCall{{ + ID: "call_heartbeat_read_file", + Type: "function", + Name: "read_file", + Arguments: map[string]any{"path": m.filePath}, + }}, + }, nil + } + + return &providers.LLMResponse{ + Content: "HEARTBEAT_OK", + ToolCalls: []providers.ToolCall{}, + }, nil +} + +func (m *toolFeedbackProvider) GetDefaultModel() string { + return "heartbeat-tool-feedback-model" +} + type toolLimitOnlyProvider struct{} func (m *toolLimitOnlyProvider) Chat( @@ -646,6 +1099,135 @@ func (m *mockCustomTool) Execute(ctx context.Context, args map[string]any) *tool return tools.SilentResult("Custom tool executed") } +type handledMediaTool struct { + store media.MediaStore + path string +} + +func (m *handledMediaTool) Name() string { return "handled_media_tool" } +func (m *handledMediaTool) Description() string { + return "Returns a media attachment and fully handles the user response" +} + +func (m *handledMediaTool) Parameters() map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{}, + } +} + +func (m *handledMediaTool) Execute(ctx context.Context, args map[string]any) *tools.ToolResult { + ref, err := m.store.Store(m.path, media.MediaMeta{ + Filename: filepath.Base(m.path), + ContentType: "image/png", + Source: "test:handled_media_tool", + }, "test:handled_media") + if err != nil { + return tools.ErrorResult(err.Error()).WithError(err) + } + return tools.MediaResult("Attachment delivered by tool.", []string{ref}).WithResponseHandled() +} + +type handledMediaWithSteeringProvider struct { + calls int +} + +func (m *handledMediaWithSteeringProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, +) (*providers.LLMResponse, error) { + m.calls++ + if m.calls == 1 { + return &providers.LLMResponse{ + Content: "Taking the screenshot now.", + ToolCalls: []providers.ToolCall{{ + ID: "call_handled_media_steering", + Type: "function", + Name: "handled_media_with_steering_tool", + Arguments: map[string]any{}, + }}, + }, nil + } + + for _, msg := range messages { + if msg.Role == "user" && msg.Content == "what about this instead?" { + return &providers.LLMResponse{Content: "Handled the queued steering message."}, nil + } + } + + return nil, fmt.Errorf("provider did not receive queued steering message") +} + +func (m *handledMediaWithSteeringProvider) GetDefaultModel() string { + return "handled-media-with-steering-model" +} + +type handledMediaWithSteeringTool struct { + store media.MediaStore + path string + loop *AgentLoop +} + +func (m *handledMediaWithSteeringTool) Name() string { return "handled_media_with_steering_tool" } +func (m *handledMediaWithSteeringTool) Description() string { + return "Returns handled media and enqueues a steering message during execution" +} + +func (m *handledMediaWithSteeringTool) Parameters() map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{}, + } +} + +func (m *handledMediaWithSteeringTool) Execute(ctx context.Context, args map[string]any) *tools.ToolResult { + if err := m.loop.Steer(providers.Message{Role: "user", Content: "what about this instead?"}); err != nil { + return tools.ErrorResult(err.Error()).WithError(err) + } + + ref, err := m.store.Store(m.path, media.MediaMeta{ + Filename: filepath.Base(m.path), + ContentType: "image/png", + Source: "test:handled_media_with_steering_tool", + }, "test:handled_media_with_steering") + if err != nil { + return tools.ErrorResult(err.Error()).WithError(err) + } + return tools.MediaResult("Attachment delivered by tool.", []string{ref}).WithResponseHandled() +} + +type mediaArtifactTool struct { + store media.MediaStore + path string +} + +func (m *mediaArtifactTool) Name() string { return "media_artifact_tool" } +func (m *mediaArtifactTool) Description() string { + return "Returns a media artifact that the agent can forward or save later" +} + +func (m *mediaArtifactTool) Parameters() map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{}, + } +} + +func (m *mediaArtifactTool) Execute(ctx context.Context, args map[string]any) *tools.ToolResult { + ref, err := m.store.Store(m.path, media.MediaMeta{ + Filename: filepath.Base(m.path), + ContentType: "image/png", + Source: "test:media_artifact_tool", + }, "test:media_artifact") + if err != nil { + return tools.ErrorResult(err.Error()).WithError(err) + } + return tools.MediaResult("Artifact created.", []string{ref}) +} + type toolLimitTestTool struct{} func (m *toolLimitTestTool) Name() string { @@ -1495,18 +2077,17 @@ func TestTargetReasoningChannelID_AllChannels(t *testing.T) { t.Fatalf("Failed to create channel manager: %v", err) } for name, id := range map[string]string{ - "whatsapp": "rid-whatsapp", - "telegram": "rid-telegram", - "feishu": "rid-feishu", - "discord": "rid-discord", - "maixcam": "rid-maixcam", - "qq": "rid-qq", - "dingtalk": "rid-dingtalk", - "slack": "rid-slack", - "line": "rid-line", - "onebot": "rid-onebot", - "wecom": "rid-wecom", - "wecom_app": "rid-wecom-app", + "whatsapp": "rid-whatsapp", + "telegram": "rid-telegram", + "feishu": "rid-feishu", + "discord": "rid-discord", + "maixcam": "rid-maixcam", + "qq": "rid-qq", + "dingtalk": "rid-dingtalk", + "slack": "rid-slack", + "line": "rid-line", + "onebot": "rid-onebot", + "wecom": "rid-wecom", } { chManager.RegisterChannel(name, &fakeChannel{id: id}) } @@ -1526,7 +2107,6 @@ func TestTargetReasoningChannelID_AllChannels(t *testing.T) { {channel: "line", wantID: "rid-line"}, {channel: "onebot", wantID: "rid-onebot"}, {channel: "wecom", wantID: "rid-wecom"}, - {channel: "wecom_app", wantID: "rid-wecom-app"}, {channel: "unknown", wantID: ""}, } @@ -1768,6 +2348,112 @@ func TestProcessMessage_PublishesReasoningContentToReasoningChannel(t *testing.T } } +func TestProcessHeartbeat_DoesNotPublishToolFeedback(t *testing.T) { + tmpDir := t.TempDir() + heartbeatFile := filepath.Join(tmpDir, "heartbeat-task.txt") + if err := os.WriteFile(heartbeatFile, []byte("heartbeat task"), 0o644); err != nil { + t.Fatalf("WriteFile() error = %v", err) + } + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + ModelName: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + ToolFeedback: config.ToolFeedbackConfig{ + Enabled: true, + MaxArgsLength: 300, + }, + }, + }, + Tools: config.ToolsConfig{ + ReadFile: config.ReadFileToolConfig{ + Enabled: true, + }, + }, + } + + msgBus := bus.NewMessageBus() + provider := &toolFeedbackProvider{filePath: heartbeatFile} + al := NewAgentLoop(cfg, msgBus, provider) + + response, err := al.ProcessHeartbeat(context.Background(), "check heartbeat tasks", "telegram", "chat-1") + if err != nil { + t.Fatalf("ProcessHeartbeat() error = %v", err) + } + if response != "HEARTBEAT_OK" { + t.Fatalf("ProcessHeartbeat() response = %q, want %q", response, "HEARTBEAT_OK") + } + + select { + case outbound := <-msgBus.OutboundChan(): + t.Fatalf("expected no outbound tool feedback during heartbeat, got %+v", outbound) + case <-time.After(200 * time.Millisecond): + } +} + +func TestProcessMessage_PublishesToolFeedbackWhenEnabled(t *testing.T) { + tmpDir := t.TempDir() + heartbeatFile := filepath.Join(tmpDir, "tool-feedback.txt") + if err := os.WriteFile(heartbeatFile, []byte("tool feedback task"), 0o644); err != nil { + t.Fatalf("WriteFile() error = %v", err) + } + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + ModelName: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + ToolFeedback: config.ToolFeedbackConfig{ + Enabled: true, + MaxArgsLength: 300, + }, + }, + }, + Tools: config.ToolsConfig{ + ReadFile: config.ReadFileToolConfig{ + Enabled: true, + }, + }, + } + + msgBus := bus.NewMessageBus() + provider := &toolFeedbackProvider{filePath: heartbeatFile} + al := NewAgentLoop(cfg, msgBus, provider) + + response, err := al.processMessage(context.Background(), bus.InboundMessage{ + Channel: "telegram", + SenderID: "user-1", + ChatID: "chat-1", + Content: "check tool feedback", + }) + if err != nil { + t.Fatalf("processMessage() error = %v", err) + } + if response != "HEARTBEAT_OK" { + t.Fatalf("processMessage() response = %q, want %q", response, "HEARTBEAT_OK") + } + + select { + case outbound := <-msgBus.OutboundChan(): + if outbound.Channel != "telegram" { + t.Fatalf("tool feedback channel = %q, want %q", outbound.Channel, "telegram") + } + if outbound.ChatID != "chat-1" { + t.Fatalf("tool feedback chatID = %q, want %q", outbound.ChatID, "chat-1") + } + if !strings.Contains(outbound.Content, "`read_file`") { + t.Fatalf("tool feedback content = %q, want read_file preview", outbound.Content) + } + case <-time.After(2 * time.Second): + t.Fatal("expected outbound tool feedback for regular messages") + } +} + func TestResolveMediaRefs_ResolvesToBase64(t *testing.T) { store := media.NewFileMediaStore() dir := t.TempDir() @@ -2132,3 +2818,111 @@ func TestFilterClientWebSearch_EmptyInput(t *testing.T) { t.Fatalf("len(result) = %d, want 0", len(result)) } } + +type overflowProvider struct { + calls int + lastMessages []providers.Message + chatFunc func(ctx context.Context, messages []providers.Message, tools []providers.ToolDefinition, model string, opts map[string]any) (*providers.LLMResponse, error) +} + +func (p *overflowProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, +) (*providers.LLMResponse, error) { + p.calls++ + p.lastMessages = append([]providers.Message(nil), messages...) + + if p.chatFunc != nil { + return p.chatFunc(ctx, messages, tools, model, opts) + } + + if p.calls == 1 { + return nil, errors.New("context_window_exceeded") + } + + return &providers.LLMResponse{ + Content: "Recovered from overflow", + }, nil +} + +func (p *overflowProvider) GetDefaultModel() string { + return "test-model" +} + +func TestProcessMessage_ContextOverflowRecovery(t *testing.T) { + al, cfg, _, _, cleanup := newTestAgentLoop(t) + defer cleanup() + _ = cfg + + provider := &overflowProvider{} + al.registry = NewAgentRegistry(al.cfg, provider) + + sessionKey := "agent:main:test-session" + agent := al.GetRegistry().GetDefaultAgent() + + for i := 0; i < 5; i++ { + agent.Sessions.AddFullMessage(sessionKey, providers.Message{Role: "user", Content: "heavy message"}) + agent.Sessions.AddFullMessage(sessionKey, providers.Message{Role: "assistant", Content: "response"}) + } + + response, err := al.processMessage(context.Background(), bus.InboundMessage{ + Channel: "test", + ChatID: "chat1", + SenderID: "user1", + SessionKey: "test-session", + Content: "trigger recovery", + }) + if err != nil { + t.Fatalf("processMessage() error = %v", err) + } + if response != "Recovered from overflow" { + t.Fatalf("response = %q, want %q", response, "Recovered from overflow") + } + + if provider.calls != 2 { + t.Fatalf("expected 2 calls, got %d", provider.calls) + } +} + +func TestProcessMessage_ContextOverflow_AnthropicStyle(t *testing.T) { + al, cfg, _, _, cleanup := newTestAgentLoop(t) + defer cleanup() + _ = cfg + + provider := &overflowProvider{} + al.registry = NewAgentRegistry(al.cfg, provider) + + recoveryMsg := "error: status 400: context_window_exceeded" + + provider.chatFunc = func( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, + ) (*providers.LLMResponse, error) { + if provider.calls == 1 { + return nil, errors.New(recoveryMsg) + } + return &providers.LLMResponse{Content: "Anthropic recovery success"}, nil + } + + response, err := al.processMessage(context.Background(), bus.InboundMessage{ + Channel: "test", + ChatID: "chat1", + SenderID: "user1", + Content: "hello", + }) + if err != nil { + t.Fatalf("processMessage() error = %v", err) + } + if !strings.Contains(response, "Anthropic recovery success") { + t.Fatalf("response = %q, want success message", response) + } + if provider.calls != 2 { + t.Fatalf("expected 2 calls for retry, got %d", provider.calls) + } +} diff --git a/pkg/channels/README.md b/pkg/channels/README.md index b7c56660b..7f238ece5 100644 --- a/pkg/channels/README.md +++ b/pkg/channels/README.md @@ -1255,8 +1255,7 @@ make test # Full test suite | `pkg/channels/onebot/` | `"onebot"` | ReactionCapable, MediaSender | | `pkg/channels/dingtalk/` | `"dingtalk"` | — | | `pkg/channels/feishu/` | `"feishu"` | — (architecture-specific build tags: `feishu_32.go` / `feishu_64.go`) | -| `pkg/channels/wecom/` | `"wecom"` | WebhookHandler, HealthChecker | -| `pkg/channels/wecom/` | `"wecom_app"` | MediaSender, WebhookHandler, HealthChecker | +| `pkg/channels/wecom/` | `"wecom"` | MediaSender | | `pkg/channels/qq/` | `"qq"` | — | | `pkg/channels/whatsapp/` | `"whatsapp"` | — (Bridge mode) | | `pkg/channels/whatsapp_native/` | `"whatsapp_native"` | — (Native whatsmeow mode) | @@ -1371,7 +1370,7 @@ agentLoop.Stop() // Stop Agent 2. **Feishu architecture-specific compilation**: The Feishu channel uses build tags to distinguish 32-bit and 64-bit architectures (`feishu_32.go` / `feishu_64.go`). Feishu uses the SDK's WebSocket mode (not HTTP webhook), so it does not implement `WebhookHandler`. -3. **WeCom has two factories**: `"wecom"` (Bot mode, webhook only) and `"wecom_app"` (App mode, supports MediaSender) are registered separately. Both implement `WebhookHandler` and `HealthChecker`. +3. **WeCom is now a single channel**: `"wecom"` is implemented as a WebSocket-based AI Bot channel with route persistence. Access control uses the shared channel allowlist mechanism. It no longer exposes the legacy webhook/app split. 4. **Pico Protocol**: `pkg/channels/pico/` implements a custom PicoClaw native protocol channel that receives messages via WebSocket webhook (`/pico/ws`). @@ -1381,4 +1380,4 @@ agentLoop.Stop() // Stop Agent 7. **PlaceholderConfig vs implementation**: `PlaceholderConfig` appears in 6 channel configs (Telegram, Discord, Slack, LINE, OneBot, Pico), but only channels that implement both `PlaceholderCapable` + `MessageEditor` (Telegram, Discord, Pico) can actually use placeholder message editing. The rest are reserved fields. -8. **ReasoningChannelID**: Most channel configs include a `reasoning_channel_id` field to route LLM reasoning/thinking output to a designated channel (WhatsApp, Telegram, Feishu, Discord, MaixCam, QQ, DingTalk, Slack, LINE, OneBot, WeCom, WeComApp). Note: `PicoConfig` does not currently expose this field. `BaseChannel` exposes this via the `WithReasoningChannelID` option and `ReasoningChannelID()` method. \ No newline at end of file +8. **ReasoningChannelID**: Most channel configs include a `reasoning_channel_id` field to route LLM reasoning/thinking output to a designated channel (WhatsApp, Telegram, Feishu, Discord, MaixCam, QQ, DingTalk, Slack, LINE, OneBot, WeCom). Note: `PicoConfig` does not currently expose this field. `BaseChannel` exposes this via the `WithReasoningChannelID` option and `ReasoningChannelID()` method. diff --git a/pkg/channels/README.zh.md b/pkg/channels/README.zh.md index 2c5e7356e..8bc8c8dbc 100644 --- a/pkg/channels/README.zh.md +++ b/pkg/channels/README.zh.md @@ -1254,8 +1254,7 @@ make test # 全量测试 | `pkg/channels/onebot/` | `"onebot"` | ReactionCapable, MediaSender | | `pkg/channels/dingtalk/` | `"dingtalk"` | — | | `pkg/channels/feishu/` | `"feishu"` | — (架构特定 build tags: `feishu_32.go` / `feishu_64.go`) | -| `pkg/channels/wecom/` | `"wecom"` | WebhookHandler, HealthChecker | -| `pkg/channels/wecom/` | `"wecom_app"` | MediaSender, WebhookHandler, HealthChecker | +| `pkg/channels/wecom/` | `"wecom"` | MediaSender | | `pkg/channels/qq/` | `"qq"` | — | | `pkg/channels/whatsapp/` | `"whatsapp"` | — (Bridge 模式) | | `pkg/channels/whatsapp_native/` | `"whatsapp_native"` | — (原生 whatsmeow 模式) | @@ -1370,7 +1369,7 @@ agentLoop.Stop() // 停止 Agent 2. **Feishu 架构特定编译**：Feishu channel 使用 build tags 区分 32 位和 64 位架构（`feishu_32.go` / `feishu_64.go`）。Feishu 使用 SDK 的 WebSocket 模式（非 HTTP webhook），因此不实现 `WebhookHandler`。 -3. **WeCom 有两个工厂**：`"wecom"`（Bot 模式，纯 webhook）和 `"wecom_app"`（应用模式，支持 MediaSender）分别注册。两者都实现了 `WebhookHandler` 和 `HealthChecker`。 +3. **WeCom 现在只有一个 channel**：`"wecom"` 采用 WebSocket AI Bot 实现，带路由持久化；访问控制走统一的 channel 白名单机制，不再保留旧的 webhook/app 双分支。 4. **Pico Protocol**：`pkg/channels/pico/` 实现了一个自定义的 PicoClaw 原生协议 channel，通过 WebSocket webhook (`/pico/ws`) 接收消息。 @@ -1380,4 +1379,4 @@ agentLoop.Stop() // 停止 Agent 7. **PlaceholderConfig 的配置与实现**：`PlaceholderConfig` 出现在 6 个 channel config 中（Telegram、Discord、Slack、LINE、OneBot、Pico），但只有实现了 `PlaceholderCapable` + `MessageEditor` 的 channel（Telegram、Discord、Pico）能真正使用占位消息编辑功能。其余 channel 的 `PlaceholderConfig` 为预留字段。 -8. **ReasoningChannelID**：大多数 channel config 都包含 `reasoning_channel_id` 字段，用于将 LLM 的思维链（reasoning/thinking）路由到指定 channel（WhatsApp、Telegram、Feishu、Discord、MaixCam、QQ、DingTalk、Slack、LINE、OneBot、WeCom、WeComApp）。注意：`PicoConfig` 目前不包含该字段。`BaseChannel` 通过 `WithReasoningChannelID` 选项和 `ReasoningChannelID()` 方法暴露此配置。 \ No newline at end of file +8. **ReasoningChannelID**：大多数 channel config 都包含 `reasoning_channel_id` 字段，用于将 LLM 的思维链（reasoning/thinking）路由到指定 channel（WhatsApp、Telegram、Feishu、Discord、MaixCam、QQ、DingTalk、Slack、LINE、OneBot、WeCom）。注意：`PicoConfig` 目前不包含该字段。`BaseChannel` 通过 `WithReasoningChannelID` 选项和 `ReasoningChannelID()` 方法暴露此配置。 diff --git a/pkg/channels/discord/discord.go b/pkg/channels/discord/discord.go index 3b5b4f8bb..2385544a6 100644 --- a/pkg/channels/discord/discord.go +++ b/pkg/channels/discord/discord.go @@ -254,10 +254,7 @@ func (c *DiscordChannel) SendPlaceholder(ctx context.Context, chatID string) (st return "", nil } - text := c.config.Placeholder.Text - if text == "" { - text = "Thinking... 💭" - } + text := c.config.Placeholder.GetRandomText() msg, err := c.session.ChannelMessageSend(chatID, text) if err != nil { diff --git a/pkg/channels/feishu/feishu_64.go b/pkg/channels/feishu/feishu_64.go index 0ab70649f..76df988ad 100644 --- a/pkg/channels/feishu/feishu_64.go +++ b/pkg/channels/feishu/feishu_64.go @@ -211,10 +211,7 @@ func (c *FeishuChannel) SendPlaceholder(ctx context.Context, chatID string) (str return "", nil } - text := c.config.Placeholder.Text - if text == "" { - text = "Thinking..." - } + text := c.config.Placeholder.GetRandomText() cardContent, err := buildMarkdownCard(text) if err != nil { diff --git a/pkg/channels/manager.go b/pkg/channels/manager.go index f04d989a3..7bcb933ce 100644 --- a/pkg/channels/manager.go +++ b/pkg/channels/manager.go @@ -206,6 +206,40 @@ func (m *Manager) preSend(ctx context.Context, name string, msg bus.OutboundMess return false } +// preSendMedia handles typing stop, reaction undo, and placeholder cleanup +// before sending media attachments. Unlike preSend for text messages, media +// delivery never edits the placeholder because there is no text payload to +// replace it with; it only attempts to delete the placeholder when possible. +func (m *Manager) preSendMedia(ctx context.Context, name string, msg bus.OutboundMediaMessage, ch Channel) { + key := name + ":" + msg.ChatID + + // 1. Stop typing + if v, loaded := m.typingStops.LoadAndDelete(key); loaded { + if entry, ok := v.(typingEntry); ok { + entry.stop() // idempotent, safe + } + } + + // 2. Undo reaction + if v, loaded := m.reactionUndos.LoadAndDelete(key); loaded { + if entry, ok := v.(reactionEntry); ok { + entry.undo() // idempotent, safe + } + } + + // 3. Clear any finalized stream marker for this chat before media delivery. + m.streamActive.LoadAndDelete(key) + + // 4. Delete placeholder if present. + if v, loaded := m.placeholders.LoadAndDelete(key); loaded { + if entry, ok := v.(placeholderEntry); ok && entry.id != "" { + if deleter, ok := ch.(MessageDeleter); ok { + deleter.DeleteMessage(ctx, msg.ChatID, entry.id) // best effort + } + } + } +} + func NewManager(cfg *config.Config, messageBus *bus.MessageBus, store media.MediaStore) (*Manager, error) { m := &Manager{ channels: make(map[string]Channel), @@ -371,19 +405,10 @@ func (m *Manager) initChannels(channels *config.ChannelsConfig) error { m.initChannel("onebot", "OneBot") } - if channels.WeCom.Enabled && channels.WeCom.Token() != "" { + if channels.WeCom.Enabled && channels.WeCom.BotID != "" && channels.WeCom.Secret() != "" { m.initChannel("wecom", "WeCom") } - if channels.WeComAIBot.Enabled && (channels.WeComAIBot.Token() != "" || - (channels.WeComAIBot.Secret() != "" && channels.WeComAIBot.BotID != "")) { - m.initChannel("wecom_aibot", "WeCom AI Bot") - } - - if channels.WeComApp.Enabled && channels.WeComApp.CorpID != "" { - m.initChannel("wecom_app", "WeCom App") - } - if channels.Weixin.Enabled && channels.Weixin.Token() != "" { m.initChannel("weixin", "Weixin") } @@ -583,8 +608,10 @@ func newChannelWorker(name string, ch Channel) *channelWorker { } } -// runWorker processes outbound messages for a single channel, splitting -// messages that exceed the channel's maximum message length. +// runWorker processes outbound messages for a single channel. +// Message processing follows this order: +// 1. SplitByMarker (if enabled in config) - LLM semantic marker-based splitting +// 2. SplitMessage - channel-specific length-based splitting (MaxMessageLength) func (m *Manager) runWorker(ctx context.Context, name string, w *channelWorker) { defer close(w.done) for { @@ -597,15 +624,29 @@ func (m *Manager) runWorker(ctx context.Context, name string, w *channelWorker) if mlp, ok := w.ch.(MessageLengthProvider); ok { maxLen = mlp.MaxMessageLength() } - if maxLen > 0 && len([]rune(msg.Content)) > maxLen { - chunks := SplitMessage(msg.Content, maxLen) - for _, chunk := range chunks { - chunkMsg := msg - chunkMsg.Content = chunk - m.sendWithRetry(ctx, name, w, chunkMsg) + + // Collect all message chunks to send + var chunks []string + + // Step 1: Try marker-based splitting if enabled + if m.config != nil && m.config.Agents.Defaults.SplitOnMarker { + if markerChunks := SplitByMarker(msg.Content); len(markerChunks) > 1 { + for _, chunk := range markerChunks { + chunks = append(chunks, splitByLength(chunk, maxLen)...) + } } - } else { - m.sendWithRetry(ctx, name, w, msg) + } + + // Step 2: Fallback to length-based splitting if no chunks from marker + if len(chunks) == 0 { + chunks = splitByLength(msg.Content, maxLen) + } + + // Step 3: Send all chunks + for _, chunk := range chunks { + chunkMsg := msg + chunkMsg.Content = chunk + m.sendWithRetry(ctx, name, w, chunkMsg) } case <-ctx.Done(): return @@ -613,6 +654,14 @@ func (m *Manager) runWorker(ctx context.Context, name string, w *channelWorker) } } +// splitByLength splits content by maxLen if needed, otherwise returns single chunk. +func splitByLength(content string, maxLen int) []string { + if maxLen > 0 && len([]rune(content)) > maxLen { + return SplitMessage(content, maxLen) + } + return []string{content} +} + // sendWithRetry sends a message through the channel with rate limiting and // retry logic. It classifies errors to determine the retry strategy: // - ErrNotRunning / ErrSendFailed: permanent, no retry @@ -774,7 +823,7 @@ func (m *Manager) runMediaWorker(ctx context.Context, name string, w *channelWor if !ok { return } - m.sendMediaWithRetry(ctx, name, w, msg) + _ = m.sendMediaWithRetry(ctx, name, w, msg) case <-ctx.Done(): return } @@ -782,26 +831,37 @@ func (m *Manager) runMediaWorker(ctx context.Context, name string, w *channelWor } // sendMediaWithRetry sends a media message through the channel with rate limiting and -// retry logic. If the channel does not implement MediaSender, it silently skips. -func (m *Manager) sendMediaWithRetry(ctx context.Context, name string, w *channelWorker, msg bus.OutboundMediaMessage) { +// retry logic. It returns nil on success, or the last error after retries, +// including when the channel does not support MediaSender. +func (m *Manager) sendMediaWithRetry( + ctx context.Context, + name string, + w *channelWorker, + msg bus.OutboundMediaMessage, +) error { ms, ok := w.ch.(MediaSender) if !ok { - logger.DebugCF("channels", "Channel does not support MediaSender, skipping media", map[string]any{ + err := fmt.Errorf("channel %q does not support media sending", name) + logger.WarnCF("channels", "Channel does not support MediaSender", map[string]any{ "channel": name, + "error": err.Error(), }) - return + return err } // Rate limit: wait for token if err := w.limiter.Wait(ctx); err != nil { - return + return err } + // Pre-send: stop typing and clean up any placeholder before sending media. + m.preSendMedia(ctx, name, msg, w.ch) + var lastErr error for attempt := 0; attempt <= maxRetries; attempt++ { lastErr = ms.SendMedia(ctx, msg) if lastErr == nil { - return + return nil } // Permanent failures — don't retry @@ -820,7 +880,7 @@ func (m *Manager) sendMediaWithRetry(ctx context.Context, name string, w *channe case <-time.After(rateLimitDelay): continue case <-ctx.Done(): - return + return ctx.Err() } } @@ -829,7 +889,7 @@ func (m *Manager) sendMediaWithRetry(ctx context.Context, name string, w *channe select { case <-time.After(backoff): case <-ctx.Done(): - return + return ctx.Err() } } @@ -840,6 +900,7 @@ func (m *Manager) sendMediaWithRetry(ctx context.Context, name string, w *channe "error": lastErr.Error(), "retries": maxRetries, }) + return lastErr } // runTTLJanitor periodically scans the typingStops and placeholders maps @@ -1032,6 +1093,26 @@ func (m *Manager) SendMessage(ctx context.Context, msg bus.OutboundMessage) erro return nil } +// SendMedia sends outbound media synchronously through the channel worker's +// rate limiter and retry logic. It blocks until the media is delivered (or all +// retries are exhausted), which preserves ordering when later agent behavior +// depends on actual media delivery. +func (m *Manager) SendMedia(ctx context.Context, msg bus.OutboundMediaMessage) error { + m.mu.RLock() + _, exists := m.channels[msg.Channel] + w, wExists := m.workers[msg.Channel] + m.mu.RUnlock() + + if !exists { + return fmt.Errorf("channel %s not found", msg.Channel) + } + if !wExists || w == nil { + return fmt.Errorf("channel %s has no active worker", msg.Channel) + } + + return m.sendMediaWithRetry(ctx, msg.Channel, w, msg) +} + func (m *Manager) SendToChannel(ctx context.Context, channelName, chatID, content string) error { m.mu.RLock() _, exists := m.channels[channelName] diff --git a/pkg/channels/manager_channel.go b/pkg/channels/manager_channel.go index 86572e336..163218b75 100644 --- a/pkg/channels/manager_channel.go +++ b/pkg/channels/manager_channel.go @@ -49,15 +49,7 @@ func hiddenValues(key string, value map[string]any, ch config.ChannelsConfig) { value["token"] = ch.LINE.ChannelAccessToken() value["secret"] = ch.LINE.ChannelSecret() case "wecom": - value["token"] = ch.WeCom.Token() - value["key"] = ch.WeCom.EncodingAESKey() - case "wecom_app": - value["token"] = ch.WeComApp.Token() - value["secret"] = ch.WeComApp.CorpSecret() - case "wecom_aibot": - value["token"] = ch.WeComAIBot.Token() - value["key"] = ch.WeComAIBot.EncodingAESKey() - value["secret"] = ch.WeComAIBot.Secret() + value["secret"] = ch.WeCom.Secret() case "dingtalk": value["secret"] = ch.QQ.AppSecret() case "qq": @@ -156,16 +148,7 @@ func updateKeys(newcfg, old *config.ChannelsConfig) { newcfg.LINE.SetChannelSecret(old.LINE.ChannelSecret()) } if newcfg.WeCom.Enabled { - newcfg.WeCom.SetToken(old.WeCom.Token()) - newcfg.WeCom.SetEncodingAESKey(old.WeCom.EncodingAESKey()) - } - if newcfg.WeComApp.Enabled { - newcfg.WeComApp.SetToken(old.WeComApp.Token()) - newcfg.WeComApp.SetCorpSecret(old.WeComApp.CorpSecret()) - } - if newcfg.WeComAIBot.Enabled { - newcfg.WeComAIBot.SetToken(old.WeComAIBot.Token()) - newcfg.WeComAIBot.SetEncodingAESKey(old.WeComAIBot.EncodingAESKey()) + newcfg.WeCom.SetSecret(old.WeCom.Secret()) } if newcfg.DingTalk.Enabled { newcfg.DingTalk.SetClientSecret(old.DingTalk.ClientSecret()) diff --git a/pkg/channels/manager_test.go b/pkg/channels/manager_test.go index 7dfec9ebf..b4fd2ba3d 100644 --- a/pkg/channels/manager_test.go +++ b/pkg/channels/manager_test.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "strings" "sync" "sync/atomic" "testing" @@ -43,6 +44,40 @@ func (m *mockChannel) EditMessage(ctx context.Context, chatID, messageID, conten return nil } +type mockMediaChannel struct { + mockChannel + sendMediaFn func(ctx context.Context, msg bus.OutboundMediaMessage) error + sentMediaMessages []bus.OutboundMediaMessage +} + +func (m *mockMediaChannel) SendMedia(ctx context.Context, msg bus.OutboundMediaMessage) error { + m.sentMediaMessages = append(m.sentMediaMessages, msg) + if m.sendMediaFn != nil { + return m.sendMediaFn(ctx, msg) + } + return nil +} + +type mockDeletingMediaChannel struct { + mockMediaChannel + deleteCalls int + lastDeleted struct { + chatID string + messageID string + } +} + +func (m *mockDeletingMediaChannel) DeleteMessage( + _ context.Context, + chatID string, + messageID string, +) error { + m.deleteCalls++ + m.lastDeleted.chatID = chatID + m.lastDeleted.messageID = messageID + return nil +} + // newTestManager creates a minimal Manager suitable for unit tests. func newTestManager() *Manager { return &Manager{ @@ -208,6 +243,125 @@ func TestSendWithRetry_MaxRetriesExhausted(t *testing.T) { } } +func TestSendMedia_Success(t *testing.T) { + m := newTestManager() + var callCount int + ch := &mockMediaChannel{ + sendMediaFn: func(_ context.Context, _ bus.OutboundMediaMessage) error { + callCount++ + return nil + }, + } + w := &channelWorker{ + ch: ch, + limiter: rate.NewLimiter(rate.Inf, 1), + } + m.channels["test"] = ch + m.workers["test"] = w + + err := m.SendMedia(context.Background(), bus.OutboundMediaMessage{ + Channel: "test", + ChatID: "chat1", + Parts: []bus.MediaPart{{Ref: "media://abc"}}, + }) + if err != nil { + t.Fatalf("SendMedia() error = %v", err) + } + if callCount != 1 { + t.Fatalf("expected 1 SendMedia call, got %d", callCount) + } +} + +func TestSendMedia_PropagatesFailure(t *testing.T) { + m := newTestManager() + ch := &mockMediaChannel{ + sendMediaFn: func(_ context.Context, _ bus.OutboundMediaMessage) error { + return fmt.Errorf("bad upload: %w", ErrSendFailed) + }, + } + w := &channelWorker{ + ch: ch, + limiter: rate.NewLimiter(rate.Inf, 1), + } + m.channels["test"] = ch + m.workers["test"] = w + + err := m.SendMedia(context.Background(), bus.OutboundMediaMessage{ + Channel: "test", + ChatID: "chat1", + Parts: []bus.MediaPart{{Ref: "media://abc"}}, + }) + if err == nil { + t.Fatal("expected SendMedia to return error") + } + if !errors.Is(err, ErrSendFailed) { + t.Fatalf("expected ErrSendFailed, got %v", err) + } +} + +func TestSendMedia_UnsupportedChannelReturnsError(t *testing.T) { + m := newTestManager() + ch := &mockChannel{ + sendFn: func(_ context.Context, _ bus.OutboundMessage) error { + return nil + }, + } + w := &channelWorker{ + ch: ch, + limiter: rate.NewLimiter(rate.Inf, 1), + } + m.channels["test"] = ch + m.workers["test"] = w + + err := m.SendMedia(context.Background(), bus.OutboundMediaMessage{ + Channel: "test", + ChatID: "chat1", + Parts: []bus.MediaPart{{Ref: "media://abc"}}, + }) + if err == nil { + t.Fatal("expected SendMedia to return error for unsupported channel") + } + if !strings.Contains(err.Error(), "does not support media sending") { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestSendMedia_DeletesPlaceholderBeforeSending(t *testing.T) { + m := newTestManager() + ch := &mockDeletingMediaChannel{ + mockMediaChannel: mockMediaChannel{ + sendMediaFn: func(_ context.Context, _ bus.OutboundMediaMessage) error { + return nil + }, + }, + } + w := &channelWorker{ + ch: ch, + limiter: rate.NewLimiter(rate.Inf, 1), + } + m.channels["test"] = ch + m.workers["test"] = w + m.RecordPlaceholder("test", "chat1", "placeholder-1") + + err := m.SendMedia(context.Background(), bus.OutboundMediaMessage{ + Channel: "test", + ChatID: "chat1", + Parts: []bus.MediaPart{{Ref: "media://abc"}}, + }) + if err != nil { + t.Fatalf("SendMedia() error = %v", err) + } + if ch.deleteCalls != 1 { + t.Fatalf("expected placeholder delete to be called once, got %d", ch.deleteCalls) + } + if ch.lastDeleted.chatID != "chat1" || ch.lastDeleted.messageID != "placeholder-1" { + t.Fatalf("unexpected placeholder deletion target: %+v", ch.lastDeleted) + } + if len(ch.sentMediaMessages) != 1 { + t.Fatalf("expected media to be sent once, got %d", len(ch.sentMediaMessages)) + } +} + func TestSendWithRetry_UnknownError(t *testing.T) { m := newTestManager() var callCount int diff --git a/pkg/channels/marker.go b/pkg/channels/marker.go new file mode 100644 index 000000000..4801e3d27 --- /dev/null +++ b/pkg/channels/marker.go @@ -0,0 +1,37 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// Inspired by and based on nanobot: https://github.com/HKUDS/nanobot +// License: MIT +// +// Copyright (c) 2026 PicoClaw contributors + +package channels + +import ( + "strings" +) + +// MessageSplitMarker is the delimiter used to split a message into multiple outbound messages. +// When SplitOnMarker is enabled in config, the Manager will split messages on this marker +// and send each part as a separate message. +const MessageSplitMarker = "<|[SPLIT]|>" + +// SplitByMarker splits a message by the MessageSplitMarker and returns the parts. +// Empty parts (including from consecutive markers) are filtered out. +// If no marker is found, returns a single-element slice containing the original content. +func SplitByMarker(content string) []string { + if content == "" { + return nil + } + parts := strings.Split(content, MessageSplitMarker) + result := make([]string, 0, len(parts)) + for _, part := range parts { + trimmed := strings.TrimSpace(part) + if trimmed != "" { + result = append(result, trimmed) + } + } + if len(result) == 0 { + return []string{content} + } + return result +} diff --git a/pkg/channels/marker_test.go b/pkg/channels/marker_test.go new file mode 100644 index 000000000..b7b4ca99e --- /dev/null +++ b/pkg/channels/marker_test.go @@ -0,0 +1,141 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// License: MIT +// +// Copyright (c) 2026 PicoClaw contributors + +package channels + +import ( + "testing" +) + +func TestSplitByMarker_Basic(t *testing.T) { + content := "Hello <|[SPLIT]|>World" + chunks := SplitByMarker(content) + + if len(chunks) != 2 { + t.Fatalf("Expected 2 chunks, got %d: %q", len(chunks), chunks) + } + if chunks[0] != "Hello" { + t.Errorf("Expected first chunk 'Hello', got %q", chunks[0]) + } + if chunks[1] != "World" { + t.Errorf("Expected second chunk 'World', got %q", chunks[1]) + } +} + +func TestSplitByMarker_NoMarker(t *testing.T) { + content := "Hello World" + chunks := SplitByMarker(content) + + if len(chunks) != 1 { + t.Fatalf("Expected 1 chunk, got %d: %q", len(chunks), chunks) + } + if chunks[0] != "Hello World" { + t.Errorf("Expected chunk 'Hello World', got %q", chunks[0]) + } +} + +func TestSplitByMarker_MultipleMarkers(t *testing.T) { + content := "Part1 <|[SPLIT]|> Part2 <|[SPLIT]|> Part3" + chunks := SplitByMarker(content) + + if len(chunks) != 3 { + t.Fatalf("Expected 3 chunks, got %d: %q", len(chunks), chunks) + } + if chunks[0] != "Part1" || chunks[1] != "Part2" || chunks[2] != "Part3" { + t.Errorf("Unexpected chunks: %q", chunks) + } +} + +func TestSplitByMarker_EmptyParts(t *testing.T) { + // Test consecutive markers and leading/trailing markers + content := "<|[SPLIT]|>Hello <|[SPLIT]|><|[SPLIT]|>World<|[SPLIT]|>" + chunks := SplitByMarker(content) + + if len(chunks) != 2 { + t.Fatalf("Expected 2 chunks, got %d: %q", len(chunks), chunks) + } + if chunks[0] != "Hello" || chunks[1] != "World" { + t.Errorf("Unexpected chunks: %q", chunks) + } +} + +func TestSplitByMarker_WhitespaceTrimmed(t *testing.T) { + content := " Hello <|[SPLIT]|> World " + chunks := SplitByMarker(content) + + if len(chunks) != 2 { + t.Fatalf("Expected 2 chunks, got %d: %q", len(chunks), chunks) + } + if chunks[0] != "Hello" || chunks[1] != "World" { + t.Errorf("Whitespace should be trimmed: %q", chunks) + } +} + +func TestSplitByMarker_EmptyInput(t *testing.T) { + chunks := SplitByMarker("") + if len(chunks) != 0 { + t.Errorf("Expected empty slice for empty input, got %d chunks", len(chunks)) + } +} + +// TestMarkerAndLengthSplitIntegration tests that SplitByMarker and SplitMessage work together correctly. +// Marker splitting happens first (per-agent config), then length splitting happens (per-channel config). +func TestMarkerAndLengthSplitIntegration(t *testing.T) { + maxLen := 10 + + // Original content: "Short <|[SPLIT]|> ThisIsAVeryLongString" + content := "Short <|[SPLIT]|> ThisIsAVeryLongString" + markerChunks := SplitByMarker(content) + + // Step 1: Marker split should give us 2 chunks + if len(markerChunks) != 2 { + t.Fatalf("Expected 2 marker chunks, got %d: %q", len(markerChunks), markerChunks) + } + + // Step 2: Length split should be applied to each marker chunk + var finalChunks []string + for _, chunk := range markerChunks { + if len([]rune(chunk)) > maxLen { + lengthChunks := SplitMessage(chunk, maxLen) + finalChunks = append(finalChunks, lengthChunks...) + } else { + finalChunks = append(finalChunks, chunk) + } + } + + // "Short" is 6 chars, within limit + // "ThisIsAVeryLongString" is 22 chars, should be split into multiple chunks + // SplitMessage with maxLen=10 splits: "ThisIsAVeryLongString" -> ["ThisI", "sAVer", "yLong", "String"] (5 chunks) + if len(finalChunks) != 5 { + t.Errorf("Expected 5 final chunks, got %d: %q", len(finalChunks), finalChunks) + } + + // Verify first chunk is unchanged + if finalChunks[0] != "Short" { + t.Errorf("First chunk should be 'Short', got %q", finalChunks[0]) + } + + // Verify all length-split chunks are within limit + for i, chunk := range finalChunks[1:] { + if len([]rune(chunk)) > maxLen { + t.Errorf("Chunk %d exceeds maxLen: %q (%d chars)", i+1, chunk, len([]rune(chunk))) + } + } +} + +// TestMarkerSplitPreservesCodeBlockIntegrity tests that marker split preserves code block boundaries +func TestMarkerSplitPreservesCodeBlockIntegrity(t *testing.T) { + content := "Hello <|[SPLIT]|>```go\npackage main\n```<|[SPLIT]|>World" + chunks := SplitByMarker(content) + + if len(chunks) != 3 { + t.Fatalf("Expected 3 chunks, got %d: %q", len(chunks), chunks) + } + + // Verify code block is intact in middle chunk + if chunks[1] != "```go\npackage main\n```" { + t.Errorf("Code block not preserved correctly: %q", chunks[1]) + } +} diff --git a/pkg/channels/matrix/init.go b/pkg/channels/matrix/init.go index 6677f855e..4d6ad45a7 100644 --- a/pkg/channels/matrix/init.go +++ b/pkg/channels/matrix/init.go @@ -1,6 +1,8 @@ package matrix import ( + "path/filepath" + "github.com/sipeed/picoclaw/pkg/bus" "github.com/sipeed/picoclaw/pkg/channels" "github.com/sipeed/picoclaw/pkg/config" @@ -8,6 +10,11 @@ import ( func init() { channels.RegisterFactory("matrix", func(cfg *config.Config, b *bus.MessageBus) (channels.Channel, error) { - return NewMatrixChannel(cfg.Channels.Matrix, b) + matrixCfg := cfg.Channels.Matrix + cryptoDatabasePath := matrixCfg.CryptoDatabasePath + if cryptoDatabasePath == "" { + cryptoDatabasePath = filepath.Join(cfg.WorkspacePath(), "matrix") + } + return NewMatrixChannel(matrixCfg, b, cryptoDatabasePath) }) } diff --git a/pkg/channels/matrix/matrix.go b/pkg/channels/matrix/matrix.go index 98c607d0b..f6370fa20 100644 --- a/pkg/channels/matrix/matrix.go +++ b/pkg/channels/matrix/matrix.go @@ -2,6 +2,7 @@ package matrix import ( "context" + "database/sql" "fmt" "html" "io" @@ -17,9 +18,12 @@ import ( "github.com/gomarkdown/markdown" mdhtml "github.com/gomarkdown/markdown/html" "github.com/gomarkdown/markdown/parser" + "go.mau.fi/util/dbutil" "maunium.net/go/mautrix" + "maunium.net/go/mautrix/crypto/cryptohelper" "maunium.net/go/mautrix/event" "maunium.net/go/mautrix/id" + _ "modernc.org/sqlite" "github.com/sipeed/picoclaw/pkg/bus" "github.com/sipeed/picoclaw/pkg/channels" @@ -30,6 +34,9 @@ import ( ) const ( + sqliteDriver = "sqlite" + dbName = "store.db" + typingRefreshInterval = 20 * time.Second typingServerTTL = 30 * time.Second roomKindCacheTTL = 5 * time.Minute @@ -181,9 +188,16 @@ type MatrixChannel struct { roomKindCache *roomKindCache localpartMentionR *regexp.Regexp + + cryptoHelper *cryptohelper.CryptoHelper + cryptoDbPath string } -func NewMatrixChannel(cfg config.MatrixConfig, messageBus *bus.MessageBus) (*MatrixChannel, error) { +func NewMatrixChannel( + cfg config.MatrixConfig, + messageBus *bus.MessageBus, + cryptoDatabasePath string, +) (*MatrixChannel, error) { homeserver := strings.TrimSpace(cfg.Homeserver) userID := strings.TrimSpace(cfg.UserID) accessToken := strings.TrimSpace(cfg.AccessToken()) @@ -230,6 +244,7 @@ func NewMatrixChannel(cfg config.MatrixConfig, messageBus *bus.MessageBus) (*Mat roomKindCache: newRoomKindCache(roomKindCacheMaxEntries, roomKindCacheTTL), localpartMentionR: localpartMentionRegexp(matrixLocalpart(client.UserID)), typingMu: sync.Mutex{}, + cryptoDbPath: cryptoDatabasePath, }, nil } @@ -239,7 +254,21 @@ func (c *MatrixChannel) Start(ctx context.Context) error { c.ctx, c.cancel = context.WithCancel(ctx) c.startTime = time.Now() + // Initialize crypto helper if database and passphrase are configured + if c.cryptoDbPath != "" && c.config.CryptoPassphrase != "" { + if err := c.initCrypto(ctx); err != nil { + logger.WarnCF( + "matrix", + "Failed to initialize crypto, continuing without encryption support", + map[string]any{ + "error": err.Error(), + }, + ) + } + } + c.syncer.OnEventType(event.EventMessage, c.handleMessageEvent) + c.syncer.OnEventType(event.EventEncrypted, c.handleMessageEvent) c.syncer.OnEventType(event.StateMember, c.handleMemberEvent) c.SetRunning(true) @@ -266,10 +295,84 @@ func (c *MatrixChannel) Stop(ctx context.Context) error { } c.stopTypingSessions(ctx) + // Close crypto helper if initialized + if c.cryptoHelper != nil { + c.cryptoHelper.Close() + c.cryptoHelper = nil + c.client.Crypto = nil + } + logger.InfoC("matrix", "Matrix channel stopped") return nil } +func (c *MatrixChannel) initCrypto(ctx context.Context) error { + logger.InfoC("matrix", "Initializing crypto helper") + + // Ensure the crypto database directory exists + if err := os.MkdirAll(c.cryptoDbPath, 0o700); err != nil { + return fmt.Errorf("create crypto database directory: %w", err) + } + + // Create database with sqlite driver (modernc.org/sqlite) + dbPath := filepath.Join(c.cryptoDbPath, dbName) + connStr := "file:" + dbPath + "?_foreign_keys=on" + + db, err := sql.Open(sqliteDriver, connStr) + if err != nil { + return fmt.Errorf("open crypto database: %w", err) + } + db.SetMaxOpenConns(1) + db.SetMaxIdleConns(1) + + // Execute PRAGMA statements + // This is equivalent to the "sqlite3-fk-wal" dialect used by cryptohelper + pragmaStmts := []string{ + "PRAGMA foreign_keys = ON", + "PRAGMA journal_mode = WAL", + "PRAGMA synchronous = NORMAL", + "PRAGMA busy_timeout = 5000", + } + for _, pragma := range pragmaStmts { + if _, err = db.ExecContext(ctx, pragma); err != nil { + _ = db.Close() + return fmt.Errorf("execute %s: %w", pragma, err) + } + } + + // Wrap with dbutil for dialect support + wrappedDB, err := dbutil.NewWithDB(db, sqliteDriver) + if err != nil { + _ = db.Close() + return fmt.Errorf("wrap database: %w", err) + } + + cryptoHelper, err := cryptohelper.NewCryptoHelper(c.client, []byte(c.config.CryptoPassphrase), wrappedDB) + if err != nil { + return fmt.Errorf("create crypto helper: %w", err) + } + + if c.client.DeviceID == "" { + resp, whoamiErr := c.client.Whoami(ctx) + if whoamiErr != nil { + _ = db.Close() + return fmt.Errorf("get device ID via whoami: %w", whoamiErr) + } + c.client.DeviceID = resp.DeviceID + } + + if err = cryptoHelper.Init(ctx); err != nil { + cryptoHelper.Close() + return fmt.Errorf("init crypto helper: %w", err) + } + + c.client.Crypto = cryptoHelper + c.cryptoHelper = cryptoHelper + + logger.InfoC("matrix", "Crypto helper initialized successfully") + return nil +} + func markdownToHTML(md string) string { p := parser.NewWithExtensions(parser.CommonExtensions | parser.AutoHeadingIDs) renderer := mdhtml.NewRenderer(mdhtml.RendererOptions{Flags: mdhtml.CommonFlags}) @@ -470,10 +573,7 @@ func (c *MatrixChannel) SendPlaceholder(ctx context.Context, chatID string) (str return "", fmt.Errorf("matrix room ID is empty") } - text := strings.TrimSpace(c.config.Placeholder.Text) - if text == "" { - text = "Thinking... 💭" - } + text := c.config.Placeholder.GetRandomText() resp, err := c.client.SendMessageEvent(ctx, roomID, event.EventMessage, &event.MessageEventContent{ MsgType: event.MsgNotice, @@ -548,9 +648,26 @@ func (c *MatrixChannel) handleMessageEvent(ctx context.Context, evt *event.Event return } - msgEvt := evt.Content.AsMessage() - if msgEvt == nil { - return + var msgEvt *event.MessageEventContent + switch evt.Type { + case event.EventMessage: + // When crypto is enabled, events marked WasEncrypted=true are + // re-dispatched by c.cryptoHelper after decryption and will be + // processed again in the EventEncrypted branch. Skip to avoid duplication. + if c.client.Crypto != nil && evt.Mautrix.WasEncrypted { + return + } + + msgEvt = evt.Content.AsMessage() + if msgEvt == nil || msgEvt.MsgType == "" { + return + } + case event.EventEncrypted: + var ok bool + msgEvt, ok = c.decryptEvent(ctx, evt) + if !ok { + return + } } // Ignore edits. @@ -642,6 +759,36 @@ func (c *MatrixChannel) handleMessageEvent(ctx context.Context, evt *event.Event ) } +// decryptEvent decrypts an encrypted event and returns the decrypted message event content. +// It returns the decrypted content and a boolean indicating whether decryption was successful. +func (c *MatrixChannel) decryptEvent(ctx context.Context, evt *event.Event) (*event.MessageEventContent, bool) { + if c.client.Crypto == nil { + logger.DebugCF("matrix", "Received encrypted message but crypto is not enabled", map[string]any{ + "room_id": evt.RoomID.String(), + }) + return nil, false + } + + decrypted, err := c.client.Crypto.Decrypt(ctx, evt) + if err != nil { + logger.WarnCF("matrix", "Failed to decrypt message", map[string]any{ + "room_id": evt.RoomID.String(), + "error": err.Error(), + }) + return nil, false + } + + if decrypted.Type != event.EventMessage { + logger.DebugCF("matrix", "Decrypted event is not a message event", map[string]any{ + "room_id": evt.RoomID.String(), + "type": decrypted.Type.String(), + }) + return nil, false + } + + return decrypted.Content.AsMessage(), true +} + func (c *MatrixChannel) extractInboundContent( ctx context.Context, msgEvt *event.MessageEventContent, diff --git a/pkg/channels/pico/pico.go b/pkg/channels/pico/pico.go index 86ce98b06..1aa1941cf 100644 --- a/pkg/channels/pico/pico.go +++ b/pkg/channels/pico/pico.go @@ -54,12 +54,13 @@ func (pc *picoConn) close() { // It serves as the reference implementation for all optional capability interfaces. type PicoChannel struct { *channels.BaseChannel - config config.PicoConfig - upgrader websocket.Upgrader - connections sync.Map // connID → *picoConn - connCount atomic.Int32 - ctx context.Context - cancel context.CancelFunc + config config.PicoConfig + upgrader websocket.Upgrader + connections map[string]*picoConn // connID -> *picoConn + sessionConnections map[string]map[string]*picoConn // sessionID -> connID -> *picoConn + connsMu sync.RWMutex + ctx context.Context + cancel context.CancelFunc } // NewPicoChannel creates a new Pico Protocol channel. @@ -92,9 +93,104 @@ func NewPicoChannel(cfg config.PicoConfig, messageBus *bus.MessageBus) (*PicoCha ReadBufferSize: 1024, WriteBufferSize: 1024, }, + connections: make(map[string]*picoConn), + sessionConnections: make(map[string]map[string]*picoConn), }, nil } +// createAndAddConnection checks MaxConnections and registers a connection atomically. +func (c *PicoChannel) createAndAddConnection(conn *websocket.Conn, sessionID string, maxConns int) (*picoConn, error) { + c.connsMu.Lock() + defer c.connsMu.Unlock() + if len(c.connections) >= maxConns { + return nil, channels.ErrTemporary + } + + var connID string + for { + connID = uuid.New().String() + if _, exists := c.connections[connID]; !exists { + break + } + } + + pc := &picoConn{ + id: connID, + conn: conn, + sessionID: sessionID, + } + + c.connections[pc.id] = pc + bySession, ok := c.sessionConnections[pc.sessionID] + if !ok { + bySession = make(map[string]*picoConn) + c.sessionConnections[pc.sessionID] = bySession + } + bySession[pc.id] = pc + + return pc, nil +} + +// removeConnection deletes a connection from indexes and returns it when found. +func (c *PicoChannel) removeConnection(connID string) *picoConn { + c.connsMu.Lock() + defer c.connsMu.Unlock() + + pc, ok := c.connections[connID] + if !ok { + return nil + } + + delete(c.connections, connID) + if bySession, ok := c.sessionConnections[pc.sessionID]; ok { + delete(bySession, connID) + if len(bySession) == 0 { + delete(c.sessionConnections, pc.sessionID) + } + } + + return pc +} + +// takeAllConnections snapshots and clears all connection indexes. +func (c *PicoChannel) takeAllConnections() []*picoConn { + c.connsMu.Lock() + defer c.connsMu.Unlock() + + all := make([]*picoConn, 0, len(c.connections)) + for _, pc := range c.connections { + all = append(all, pc) + } + clear(c.connections) + clear(c.sessionConnections) + + return all +} + +// sessionConnectionsSnapshot returns all active connections for a session. +func (c *PicoChannel) sessionConnectionsSnapshot(sessionID string) []*picoConn { + c.connsMu.RLock() + defer c.connsMu.RUnlock() + + bySession, ok := c.sessionConnections[sessionID] + if !ok || len(bySession) == 0 { + return nil + } + + conns := make([]*picoConn, 0, len(bySession)) + for _, pc := range bySession { + conns = append(conns, pc) + } + return conns +} + +// currentConnCount returns a lock-protected snapshot of active connection count. +func (c *PicoChannel) currentConnCount() int { + c.connsMu.RLock() + defer c.connsMu.RUnlock() + return len(c.connections) +} + // Start implements Channel. func (c *PicoChannel) Start(ctx context.Context) error { logger.InfoC("pico", "Starting Pico Protocol channel") @@ -110,13 +206,9 @@ func (c *PicoChannel) Stop(ctx context.Context) error { c.SetRunning(false) // Close all connections - c.connections.Range(func(key, value any) bool { - if pc, ok := value.(*picoConn); ok { - pc.close() - } - c.connections.Delete(key) - return true - }) + for _, pc := range c.takeAllConnections() { + pc.close() + } if c.cancel != nil { c.cancel() @@ -133,8 +225,8 @@ func (c *PicoChannel) WebhookPath() string { return "/pico/" } func (c *PicoChannel) ServeHTTP(w http.ResponseWriter, r *http.Request) { path := strings.TrimPrefix(r.URL.Path, "/pico") - switch { - case path == "/ws" || path == "/ws/": + switch path { + case "/ws", "/ws/": c.handleWebSocket(w, r) default: http.NotFound(w, r) @@ -183,10 +275,7 @@ func (c *PicoChannel) SendPlaceholder(ctx context.Context, chatID string) (strin return "", nil } - text := c.config.Placeholder.Text - if text == "" { - text = "Thinking... 💭" - } + text := c.config.Placeholder.GetRandomText() msgID := uuid.New().String() outMsg := newMessage(TypeMessageCreate, map[string]any{ @@ -208,23 +297,16 @@ func (c *PicoChannel) broadcastToSession(chatID string, msg PicoMessage) error { msg.SessionID = sessionID var sent bool - c.connections.Range(func(key, value any) bool { - pc, ok := value.(*picoConn) - if !ok { - return true + for _, pc := range c.sessionConnectionsSnapshot(sessionID) { + if err := pc.writeJSON(msg); err != nil { + logger.DebugCF("pico", "Write to connection failed", map[string]any{ + "conn_id": pc.id, + "error": err.Error(), + }) + } else { + sent = true } - if pc.sessionID == sessionID { - if err := pc.writeJSON(msg); err != nil { - logger.DebugCF("pico", "Write to connection failed", map[string]any{ - "conn_id": pc.id, - "error": err.Error(), - }) - } else { - sent = true - } - } - return true - }) + } if !sent { return fmt.Errorf("no active connections for session %s: %w", sessionID, channels.ErrSendFailed) @@ -250,7 +332,7 @@ func (c *PicoChannel) handleWebSocket(w http.ResponseWriter, r *http.Request) { if maxConns <= 0 { maxConns = 100 } - if int(c.connCount.Load()) >= maxConns { + if c.currentConnCount() >= maxConns { http.Error(w, "too many connections", http.StatusServiceUnavailable) return } @@ -275,15 +357,17 @@ func (c *PicoChannel) handleWebSocket(w http.ResponseWriter, r *http.Request) { sessionID = uuid.New().String() } - pc := &picoConn{ - id: uuid.New().String(), - conn: conn, - sessionID: sessionID, + pc, err := c.createAndAddConnection(conn, sessionID, maxConns) + if err != nil { + _ = conn.WriteControl( + websocket.CloseMessage, + websocket.FormatCloseMessage(websocket.CloseTryAgainLater, "too many connections"), + time.Now().Add(2*time.Second), + ) + _ = conn.Close() + return } - c.connections.Store(pc.id, pc) - c.connCount.Add(1) - logger.InfoCF("pico", "WebSocket client connected", map[string]any{ "conn_id": pc.id, "session_id": sessionID, @@ -341,12 +425,12 @@ func (c *PicoChannel) matchedSubprotocol(r *http.Request) string { func (c *PicoChannel) readLoop(pc *picoConn) { defer func() { pc.close() - c.connections.Delete(pc.id) - c.connCount.Add(-1) - logger.InfoCF("pico", "WebSocket client disconnected", map[string]any{ - "conn_id": pc.id, - "session_id": pc.sessionID, - }) + if removed := c.removeConnection(pc.id); removed != nil { + logger.InfoCF("pico", "WebSocket client disconnected", map[string]any{ + "conn_id": removed.id, + "session_id": removed.sessionID, + }) + } }() readTimeout := time.Duration(c.config.ReadTimeout) * time.Second diff --git a/pkg/channels/pico/pico_test.go b/pkg/channels/pico/pico_test.go new file mode 100644 index 000000000..e712767ad --- /dev/null +++ b/pkg/channels/pico/pico_test.go @@ -0,0 +1,144 @@ +package pico + +import ( + "context" + "errors" + "fmt" + "sync" + "testing" + + "github.com/sipeed/picoclaw/pkg/bus" + "github.com/sipeed/picoclaw/pkg/channels" + "github.com/sipeed/picoclaw/pkg/config" +) + +func newTestPicoChannel(t *testing.T) *PicoChannel { + t.Helper() + + cfg := config.PicoConfig{} + cfg.SetToken("test-token") + ch, err := NewPicoChannel(cfg, bus.NewMessageBus()) + if err != nil { + t.Fatalf("NewPicoChannel: %v", err) + } + + ch.ctx = context.Background() + return ch +} + +func TestCreateAndAddConnection_RespectsMaxConnectionsConcurrently(t *testing.T) { + ch := newTestPicoChannel(t) + + const ( + maxConns = 5 + goroutines = 64 + sessionID = "session-a" + ) + + var wg sync.WaitGroup + var mu sync.Mutex + successCount := 0 + errCount := 0 + + wg.Add(goroutines) + for i := 0; i < goroutines; i++ { + go func() { + defer wg.Done() + + pc, err := ch.createAndAddConnection(nil, sessionID, maxConns) + mu.Lock() + defer mu.Unlock() + + if err == nil { + successCount++ + if pc == nil { + t.Errorf("pc is nil on success") + } + return + } + if !errors.Is(err, channels.ErrTemporary) { + t.Errorf("unexpected error: %v", err) + return + } + errCount++ + }() + } + wg.Wait() + + if successCount > maxConns { + t.Fatalf("successCount=%d > maxConns=%d", successCount, maxConns) + } + if successCount+errCount != goroutines { + t.Fatalf("success=%d err=%d total=%d want=%d", successCount, errCount, successCount+errCount, goroutines) + } + if got := ch.currentConnCount(); got != maxConns { + t.Fatalf("currentConnCount=%d want=%d", got, maxConns) + } +} + +func TestRemoveConnection_CleansBothIndexes(t *testing.T) { + ch := newTestPicoChannel(t) + + pc, err := ch.createAndAddConnection(nil, "session-cleanup", 10) + if err != nil { + t.Fatalf("createAndAddConnection: %v", err) + } + + removed := ch.removeConnection(pc.id) + if removed == nil { + t.Fatal("removeConnection returned nil") + } + + ch.connsMu.RLock() + defer ch.connsMu.RUnlock() + + if _, ok := ch.connections[pc.id]; ok { + t.Fatalf("connID %s still exists in connections", pc.id) + } + if _, ok := ch.sessionConnections[pc.sessionID]; ok { + t.Fatalf("session %s still exists in sessionConnections", pc.sessionID) + } + if got := len(ch.connections); got != 0 { + t.Fatalf("len(connections)=%d want=0", got) + } +} + +func TestBroadcastToSession_TargetsOnlyRequestedSession(t *testing.T) { + ch := newTestPicoChannel(t) + + target := &picoConn{id: "target", sessionID: "s-target"} + target.closed.Store(true) + ch.addConnForTest(target) + + other := &picoConn{id: "other", sessionID: "s-other"} + ch.addConnForTest(other) + + err := ch.broadcastToSession("pico:s-target", newMessage(TypeMessageCreate, map[string]any{"content": "hello"})) + if err == nil { + t.Fatal("expected send failure due to closed target connection") + } + if !errors.Is(err, channels.ErrSendFailed) { + t.Fatalf("expected ErrSendFailed, got %v", err) + } +} + +func (c *PicoChannel) addConnForTest(pc *picoConn) { + c.connsMu.Lock() + defer c.connsMu.Unlock() + if c.connections == nil { + c.connections = make(map[string]*picoConn) + } + if c.sessionConnections == nil { + c.sessionConnections = make(map[string]map[string]*picoConn) + } + if _, exists := c.connections[pc.id]; exists { + panic(fmt.Sprintf("duplicate conn id in test: %s", pc.id)) + } + c.connections[pc.id] = pc + bySession, ok := c.sessionConnections[pc.sessionID] + if !ok { + bySession = make(map[string]*picoConn) + c.sessionConnections[pc.sessionID] = bySession + } + bySession[pc.id] = pc +} diff --git a/pkg/channels/telegram/telegram.go b/pkg/channels/telegram/telegram.go index d0011d21b..5adb40a7e 100644 --- a/pkg/channels/telegram/telegram.go +++ b/pkg/channels/telegram/telegram.go @@ -402,10 +402,7 @@ func (c *TelegramChannel) SendPlaceholder(ctx context.Context, chatID string) (s return "", nil } - text := phCfg.Text - if text == "" { - text = "Thinking... 💭" - } + text := phCfg.GetRandomText() cid, threadID, err := parseTelegramChatID(chatID) if err != nil { @@ -642,8 +639,12 @@ func (c *TelegramChannel) handleMessage(ctx context.Context, message *telego.Mes } } + if content == "" && len(mediaPaths) == 0 { + return nil + } + if content == "" { - content = "[empty message]" + content = "[media only]" } // In group chats, apply unified group trigger filtering diff --git a/pkg/channels/telegram/telegram_test.go b/pkg/channels/telegram/telegram_test.go index 6bf1077af..fd189d9a7 100644 --- a/pkg/channels/telegram/telegram_test.go +++ b/pkg/channels/telegram/telegram_test.go @@ -641,3 +641,35 @@ func TestHandleMessage_ReplyThread_NonForum_NoIsolation(t *testing.T) { assert.Empty(t, inbound.Metadata["parent_peer_kind"]) assert.Empty(t, inbound.Metadata["parent_peer_id"]) } + +func TestHandleMessage_EmptyContent_Ignored(t *testing.T) { + messageBus := bus.NewMessageBus() + ch := &TelegramChannel{ + BaseChannel: channels.NewBaseChannel("telegram", nil, messageBus, nil), + chatIDs: make(map[string]int64), + ctx: context.Background(), + } + + // Service message with no text/caption/media (like ForumTopicCreated) + msg := &telego.Message{ + MessageID: 123, + Chat: telego.Chat{ + ID: 456, + Type: "group", + }, + From: &telego.User{ + ID: 789, + FirstName: "User", + }, + } + + err := ch.handleMessage(context.Background(), msg) + require.NoError(t, err) + + // Should NOT publish to message bus + select { + case <-messageBus.InboundChan(): + t.Fatal("Empty message should not be published to message bus") + default: + } +} diff --git a/pkg/channels/wecom/aibot.go b/pkg/channels/wecom/aibot.go deleted file mode 100644 index c5e148185..000000000 --- a/pkg/channels/wecom/aibot.go +++ /dev/null @@ -1,1099 +0,0 @@ -package wecom - -import ( - "bytes" - "context" - "crypto/rand" - "encoding/base64" - "encoding/json" - "fmt" - "io" - "math/big" - "net/http" - "strings" - "sync" - "time" - - "github.com/sipeed/picoclaw/pkg/bus" - "github.com/sipeed/picoclaw/pkg/channels" - "github.com/sipeed/picoclaw/pkg/config" - "github.com/sipeed/picoclaw/pkg/identity" - "github.com/sipeed/picoclaw/pkg/logger" - "github.com/sipeed/picoclaw/pkg/utils" -) - -// responseURLHTTPClient is a shared HTTP client for posting to WeCom response_url. -// Reusing it enables connection pooling across replies. -var responseURLHTTPClient = &http.Client{Timeout: 15 * time.Second} - -// WeComAIBotChannel implements the Channel interface for WeCom AI Bot (企业微信智能机器人) -type WeComAIBotChannel struct { - *channels.BaseChannel - config config.WeComAIBotConfig - ctx context.Context - cancel context.CancelFunc - streamTasks map[string]*streamTask // streamID -> task (for poll lookups) - chatTasks map[string][]*streamTask // chatID -> in-flight tasks queue (FIFO) - taskMu sync.RWMutex -} - -// streamTask represents a streaming task for AI Bot. -// -// Mutable fields (Finished, StreamClosed, StreamClosedAt) must be read/written -// while holding WeComAIBotChannel.taskMu. Immutable fields (StreamID, ChatID, -// ResponseURL, Question, CreatedTime, Deadline, answerCh, ctx, cancel) are set -// once at creation and never modified, so they are safe to read without a lock. -type streamTask struct { - // immutable after creation - StreamID string - ChatID string // used by Send() to find this task - ResponseURL string // temporary URL for proactive reply (valid 1 hour, use once) - Question string - CreatedTime time.Time - Deadline time.Time // ~30s, we close the stream here and switch to response_url - answerCh chan string // receives agent reply from Send() - ctx context.Context // canceled when task is removed; used to interrupt the agent goroutine - cancel context.CancelFunc // call on task removal to cancel ctx - - // mutable — guarded by WeComAIBotChannel.taskMu - StreamClosed bool // stream returned finish:true; waiting for agent to reply via response_url - StreamClosedAt time.Time // set when StreamClosed becomes true; used for accelerated cleanup - Finished bool // fully done -} - -// WeComAIBotMessage represents the decrypted JSON message from WeCom AI Bot -// Ref: https://developer.work.weixin.qq.com/document/path/100719 -type WeComAIBotMessage struct { - MsgID string `json:"msgid"` - AIBotID string `json:"aibotid"` - ChatID string `json:"chatid"` // only for group chat - ChatType string `json:"chattype"` // "single" or "group" - From struct { - UserID string `json:"userid"` - } `json:"from"` - ResponseURL string `json:"response_url"` // temporary URL for proactive reply - MsgType string `json:"msgtype"` - // text message - Text *struct { - Content string `json:"content"` - } `json:"text,omitempty"` - // stream polling refresh - Stream *struct { - ID string `json:"id"` - } `json:"stream,omitempty"` - // image message - Image *struct { - URL string `json:"url"` - } `json:"image,omitempty"` - // mixed message (text + image) - Mixed *struct { - MsgItem []struct { - MsgType string `json:"msgtype"` - Text *struct { - Content string `json:"content"` - } `json:"text,omitempty"` - Image *struct { - URL string `json:"url"` - } `json:"image,omitempty"` - } `json:"msg_item"` - } `json:"mixed,omitempty"` - // event field - Event *struct { - EventType string `json:"eventtype"` - } `json:"event,omitempty"` -} - -// WeComAIBotMsgItemImage holds the image payload inside a stream message item. -type WeComAIBotMsgItemImage struct { - Base64 string `json:"base64"` - MD5 string `json:"md5"` -} - -// WeComAIBotMsgItem is a single item inside a stream's msg_item list. -type WeComAIBotMsgItem struct { - MsgType string `json:"msgtype"` - Image *WeComAIBotMsgItemImage `json:"image,omitempty"` -} - -// WeComAIBotStreamInfo represents the detailed stream content in streaming responses. -type WeComAIBotStreamInfo struct { - ID string `json:"id"` - Finish bool `json:"finish"` - Content string `json:"content,omitempty"` - MsgItem []WeComAIBotMsgItem `json:"msg_item,omitempty"` -} - -// WeComAIBotStreamResponse represents the streaming response format -type WeComAIBotStreamResponse struct { - MsgType string `json:"msgtype"` - Stream WeComAIBotStreamInfo `json:"stream"` -} - -// WeComAIBotEncryptedResponse represents the encrypted response wrapper -// Fields match WXBizJsonMsgCrypt.generate() in Python SDK -type WeComAIBotEncryptedResponse struct { - Encrypt string `json:"encrypt"` - MsgSignature string `json:"msgsignature"` - Timestamp string `json:"timestamp"` - Nonce string `json:"nonce"` -} - -// NewWeComAIBotChannel creates a WeCom AI Bot channel instance. -// If cfg.BotID and cfg.secret are both set, it returns a WeComAIBotWSChannel -// using the WebSocket long-connection API. -// Otherwise it returns the webhook-mode WeComAIBotChannel (requires Token + -// EncodingAESKey). -func NewWeComAIBotChannel( - cfg config.WeComAIBotConfig, - messageBus *bus.MessageBus, -) (channels.Channel, error) { - // WebSocket long-connection mode takes priority when BotID + secret are set. - if cfg.BotID != "" && cfg.Secret() != "" { - logger.InfoC("wecom_aibot", "BotID and secret provided, using WebSocket mode") - return newWeComAIBotWSChannel(cfg, messageBus) - } - // Webhook (short-connection) mode. - if cfg.Token() == "" || cfg.EncodingAESKey() == "" { - return nil, fmt.Errorf( - "WeCom AI Bot requires either (bot_id + secret) for WebSocket mode " + - "or (token + encoding_aes_key) for webhook mode") - } - if cfg.ProcessingMessage == "" { - cfg.ProcessingMessage = config.DefaultWeComAIBotProcessingMessage - } - - base := channels.NewBaseChannel("wecom_aibot", cfg, messageBus, cfg.AllowFrom, - channels.WithMaxMessageLength(2048), - channels.WithReasoningChannelID(cfg.ReasoningChannelID), - ) - - return &WeComAIBotChannel{ - BaseChannel: base, - config: cfg, - streamTasks: make(map[string]*streamTask), - chatTasks: make(map[string][]*streamTask), - }, nil -} - -// Name returns the channel name -func (c *WeComAIBotChannel) Name() string { - return "wecom_aibot" -} - -// Start initializes the WeCom AI Bot channel -func (c *WeComAIBotChannel) Start(ctx context.Context) error { - logger.InfoC("wecom_aibot", "Starting WeCom AI Bot channel...") - - c.ctx, c.cancel = context.WithCancel(ctx) - - // Start cleanup goroutine for old tasks - go c.cleanupLoop() - - c.SetRunning(true) - logger.InfoC("wecom_aibot", "WeCom AI Bot channel started") - - return nil -} - -// Stop gracefully stops the WeCom AI Bot channel -func (c *WeComAIBotChannel) Stop(ctx context.Context) error { - logger.InfoC("wecom_aibot", "Stopping WeCom AI Bot channel...") - - if c.cancel != nil { - c.cancel() - } - - c.SetRunning(false) - logger.InfoC("wecom_aibot", "WeCom AI Bot channel stopped") - return nil -} - -// Send delivers the agent reply into the active streamTask for msg.ChatID. -// It writes into the earliest unfinished task in the queue (FIFO per chatID). -// If the stream has already closed (deadline passed), it posts directly to response_url. -func (c *WeComAIBotChannel) Send(ctx context.Context, msg bus.OutboundMessage) error { - if !c.IsRunning() { - return channels.ErrNotRunning - } - c.taskMu.Lock() - queue := c.chatTasks[msg.ChatID] - // Only compact Finished tasks at the head of the queue. - // Tasks that are Finished in the middle are NOT removed here: doing a full - // scan on every Send() call would be O(n) and is unnecessary given that - // removeTask() always splices the task out of the queue immediately. - // Any Finished task left stranded in the middle (e.g. due to an unexpected - // code path) will be collected by cleanupOldTasks. - for len(queue) > 0 && queue[0].Finished { - queue = queue[1:] - } - c.chatTasks[msg.ChatID] = queue - var task *streamTask - var streamClosed bool - var responseURL string - if len(queue) > 0 { - task = queue[0] - // Read mutable fields while holding c.taskMu to avoid data races. - streamClosed = task.StreamClosed - responseURL = task.ResponseURL - } - c.taskMu.Unlock() - - if task == nil { - logger.DebugCF( - "wecom_aibot", - "Send: no active task for chat (may have timed out)", - map[string]any{ - "chat_id": msg.ChatID, - }, - ) - return nil - } - - if streamClosed { - // Stream already ended with a "please wait" notice; send the real reply via response_url. - // Note: task.StreamID and task.ChatID are immutable, safe to read without a lock. - logger.InfoCF("wecom_aibot", "Sending reply via response_url", map[string]any{ - "stream_id": task.StreamID, - "chat_id": msg.ChatID, - }) - if responseURL != "" { - if err := c.sendViaResponseURL(responseURL, msg.Content); err != nil { - logger.ErrorCF("wecom_aibot", "Failed to send via response_url", map[string]any{ - "error": err, - "stream_id": task.StreamID, - }) - c.removeTask(task) - return fmt.Errorf("response_url delivery failed: %w", channels.ErrSendFailed) - } - } else { - logger.WarnCF("wecom_aibot", "Stream closed but no response_url available", map[string]any{ - "stream_id": task.StreamID, - }) - } - c.removeTask(task) - return nil - } - - // Stream still open: deliver via answerCh for the next poll response. - select { - case task.answerCh <- msg.Content: - case <-task.ctx.Done(): - // Task was canceled (cleanup removed it); silently drop the reply. - return nil - case <-ctx.Done(): - return ctx.Err() - } - return nil -} - -// WebhookPath returns the path for registering on the shared HTTP server -func (c *WeComAIBotChannel) WebhookPath() string { - if c.config.WebhookPath == "" { - return "/webhook/wecom-aibot" - } - return c.config.WebhookPath -} - -// ServeHTTP implements http.Handler for the shared HTTP server -func (c *WeComAIBotChannel) ServeHTTP(w http.ResponseWriter, r *http.Request) { - c.handleWebhook(w, r) -} - -// HealthPath returns the health check endpoint path -func (c *WeComAIBotChannel) HealthPath() string { - return c.WebhookPath() + "/health" -} - -// HealthHandler handles health check requests -func (c *WeComAIBotChannel) HealthHandler(w http.ResponseWriter, r *http.Request) { - c.handleHealth(w, r) -} - -// handleWebhook handles incoming webhook requests from WeCom AI Bot -func (c *WeComAIBotChannel) handleWebhook(w http.ResponseWriter, r *http.Request) { - ctx := r.Context() - - // Log all incoming requests for debugging - logger.DebugCF("wecom_aibot", "Received webhook request", map[string]any{ - "method": r.Method, - "path": r.URL.Path, - "query": r.URL.RawQuery, - }) - - switch r.Method { - case http.MethodGet: - // URL verification - c.handleVerification(ctx, w, r) - case http.MethodPost: - // Message callback - c.handleMessageCallback(ctx, w, r) - default: - http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) - } -} - -// handleVerification handles the URL verification request from WeCom -func (c *WeComAIBotChannel) handleVerification( - ctx context.Context, - w http.ResponseWriter, - r *http.Request, -) { - msgSignature := r.URL.Query().Get("msg_signature") - timestamp := r.URL.Query().Get("timestamp") - nonce := r.URL.Query().Get("nonce") - echostr := r.URL.Query().Get("echostr") - - logger.DebugCF("wecom_aibot", "URL verification request", map[string]any{ - "msg_signature": msgSignature, - "timestamp": timestamp, - "nonce": nonce, - }) - - // Verify signature - if !verifySignature(c.config.Token(), msgSignature, timestamp, nonce, echostr) { - logger.ErrorC("wecom_aibot", "Signature verification failed") - http.Error(w, "Signature verification failed", http.StatusUnauthorized) - return - } - - // Decrypt echostr - // For WeCom AI Bot (智能机器人), receiveid should be empty string - decrypted, err := decryptMessageWithVerify(echostr, c.config.EncodingAESKey(), "") - if err != nil { - logger.ErrorCF("wecom_aibot", "Failed to decrypt echostr", map[string]any{ - "error": err, - }) - http.Error(w, "Decryption failed", http.StatusInternalServerError) - return - } - - // Remove BOM and whitespace as per WeCom documentation - decrypted = strings.TrimPrefix(decrypted, "\ufeff") - decrypted = strings.TrimSpace(decrypted) - - logger.InfoC("wecom_aibot", "URL verification successful") - w.Header().Set("Content-Type", "text/plain; charset=utf-8") - w.WriteHeader(http.StatusOK) - w.Write([]byte(decrypted)) -} - -// handleMessageCallback handles incoming messages from WeCom AI Bot -func (c *WeComAIBotChannel) handleMessageCallback( - ctx context.Context, - w http.ResponseWriter, - r *http.Request, -) { - msgSignature := r.URL.Query().Get("msg_signature") - timestamp := r.URL.Query().Get("timestamp") - nonce := r.URL.Query().Get("nonce") - - // Read request body (limit to 4 MB to prevent memory exhaustion) - const maxBodySize = 4 << 20 // 4 MB - body, err := io.ReadAll(io.LimitReader(r.Body, maxBodySize+1)) - if err != nil { - logger.ErrorCF("wecom_aibot", "Failed to read request body", map[string]any{ - "error": err, - }) - http.Error(w, "Failed to read body", http.StatusBadRequest) - return - } - if len(body) > maxBodySize { - http.Error(w, "Request body too large", http.StatusRequestEntityTooLarge) - return - } - - // Parse JSON body to get encrypted message - // Format: {"encrypt": "base64_encrypted_string"} - var encryptedMsg struct { - Encrypt string `json:"encrypt"` - } - if unmarshalErr := json.Unmarshal(body, &encryptedMsg); unmarshalErr != nil { - logger.ErrorCF("wecom_aibot", "Failed to parse JSON body", map[string]any{ - "error": unmarshalErr, - "body": string(body), - }) - http.Error(w, "Failed to parse JSON", http.StatusBadRequest) - return - } - - // Verify signature - if !verifySignature(c.config.Token(), msgSignature, timestamp, nonce, encryptedMsg.Encrypt) { - logger.ErrorC("wecom_aibot", "Signature verification failed") - http.Error(w, "Signature verification failed", http.StatusUnauthorized) - return - } - - // Decrypt message - // For WeCom AI Bot (智能机器人), receiveid is empty string - decrypted, err := decryptMessageWithVerify(encryptedMsg.Encrypt, c.config.EncodingAESKey(), "") - if err != nil { - logger.ErrorCF("wecom_aibot", "Failed to decrypt message", map[string]any{ - "error": err, - }) - http.Error(w, "Decryption failed", http.StatusInternalServerError) - return - } - - // Parse decrypted JSON message - var msg WeComAIBotMessage - if unmarshalErr := json.Unmarshal([]byte(decrypted), &msg); unmarshalErr != nil { - logger.ErrorCF("wecom_aibot", "Failed to parse decrypted JSON", map[string]any{ - "error": unmarshalErr, - "decrypted": decrypted, - }) - http.Error(w, "Failed to parse message", http.StatusInternalServerError) - return - } - - logger.DebugCF("wecom_aibot", "Decrypted message", map[string]any{ - "msgtype": msg.MsgType, - }) - - // Process the message and get streaming response - response := c.processMessage(ctx, msg, timestamp, nonce) - - // Check if response is empty (e.g. due to unsupported message type) - if response == "" { - response = c.encryptEmptyResponse(timestamp, nonce) - } - - // Return encrypted JSON response - w.Header().Set("Content-Type", "application/json; charset=utf-8") - w.WriteHeader(http.StatusOK) - w.Write([]byte(response)) -} - -// processMessage processes the received message and returns encrypted response -func (c *WeComAIBotChannel) processMessage( - ctx context.Context, - msg WeComAIBotMessage, - timestamp, nonce string, -) string { - logger.DebugCF("wecom_aibot", "Processing message", map[string]any{ - "msgtype": msg.MsgType, - }) - - switch msg.MsgType { - case "text": - return c.handleTextMessage(ctx, msg, timestamp, nonce) - case "stream": - return c.handleStreamMessage(ctx, msg, timestamp, nonce) - case "image": - return c.handleImageMessage(ctx, msg, timestamp, nonce) - case "mixed": - return c.handleMixedMessage(ctx, msg, timestamp, nonce) - case "event": - return c.handleEventMessage(ctx, msg, timestamp, nonce) - default: - logger.WarnCF("wecom_aibot", "Unsupported message type", map[string]any{ - "msgtype": msg.MsgType, - }) - return c.encryptResponse("", timestamp, nonce, WeComAIBotStreamResponse{ - MsgType: "stream", - Stream: WeComAIBotStreamInfo{ - ID: c.generateStreamID(), - Finish: true, - Content: "Unsupported message type: " + msg.MsgType, - }, - }) - } -} - -// handleTextMessage handles text messages by starting a new streaming task -func (c *WeComAIBotChannel) handleTextMessage( - ctx context.Context, - msg WeComAIBotMessage, - timestamp, nonce string, -) string { - if msg.Text == nil { - logger.ErrorC("wecom_aibot", "text message missing text field") - return c.encryptEmptyResponse(timestamp, nonce) - } - - content := msg.Text.Content - userID := msg.From.UserID - if userID == "" { - userID = "unknown" - } - - // chatID: group chat uses chatid, single chat uses userid - chatID := msg.ChatID - if chatID == "" { - chatID = userID - } - - streamID := c.generateStreamID() - - // WeCom stops sending stream-refresh callbacks after 6 minutes. - // Set a slightly shorter deadline so we can send a timeout notice before it gives up. - deadline := time.Now().Add(30 * time.Second) - - // Each task gets its own context derived from the channel lifetime context. - // Canceling taskCancel interrupts the agent goroutine when the task is removed. - taskCtx, taskCancel := context.WithCancel(c.ctx) - - task := &streamTask{ - StreamID: streamID, - ChatID: chatID, - ResponseURL: msg.ResponseURL, - Question: content, - CreatedTime: time.Now(), - Deadline: deadline, - Finished: false, - answerCh: make(chan string, 1), - ctx: taskCtx, - cancel: taskCancel, - } - - c.taskMu.Lock() - c.streamTasks[streamID] = task - c.chatTasks[chatID] = append(c.chatTasks[chatID], task) - c.taskMu.Unlock() - - // Publish to agent asynchronously; agent will call Send() with reply. - // Use task.ctx (not c.ctx) so the agent goroutine is canceled when the task is removed. - go func() { - sender := bus.SenderInfo{ - Platform: "wecom_aibot", - PlatformID: userID, - CanonicalID: identity.BuildCanonicalID("wecom_aibot", userID), - DisplayName: userID, - } - peerKind := "direct" - if msg.ChatType == "group" { - peerKind = "group" - } - peer := bus.Peer{Kind: peerKind, ID: chatID} - metadata := map[string]string{ - "channel": "wecom_aibot", - "chat_type": msg.ChatType, - "msg_type": "text", - "msgid": msg.MsgID, - "aibotid": msg.AIBotID, - "stream_id": streamID, - "response_url": msg.ResponseURL, - } - c.HandleMessage(task.ctx, peer, msg.MsgID, userID, chatID, - content, nil, metadata, sender) - }() - - // Return first streaming response immediately (finish=false, content empty) - return c.getStreamResponse(task, timestamp, nonce) -} - -// handleStreamMessage handles stream polling requests -func (c *WeComAIBotChannel) handleStreamMessage( - ctx context.Context, - msg WeComAIBotMessage, - timestamp, nonce string, -) string { - if msg.Stream == nil { - logger.ErrorC("wecom_aibot", "Stream message missing stream field") - return c.encryptEmptyResponse(timestamp, nonce) - } - - streamID := msg.Stream.ID - - c.taskMu.RLock() - task, exists := c.streamTasks[streamID] - c.taskMu.RUnlock() - - if !exists { - logger.DebugCF( - "wecom_aibot", - "Stream task not found (may be from previous session)", - map[string]any{ - "stream_id": streamID, - }, - ) - return c.encryptResponse(streamID, timestamp, nonce, WeComAIBotStreamResponse{ - MsgType: "stream", - Stream: WeComAIBotStreamInfo{ - ID: streamID, - Finish: true, - Content: "Task not found or already finished. Please resend your message to start a new session.", - }, - }) - } - - // Get next response - return c.getStreamResponse(task, timestamp, nonce) -} - -// handleImageMessage handles image messages -func (c *WeComAIBotChannel) handleImageMessage( - ctx context.Context, - msg WeComAIBotMessage, - timestamp, nonce string, -) string { - logger.WarnC("wecom_aibot", "Image message type not yet fully implemented") - if msg.Image == nil { - logger.ErrorC("wecom_aibot", "Image message missing image field") - return c.encryptEmptyResponse(timestamp, nonce) - } - - imageURL := msg.Image.URL - - // For now, just acknowledge receipt without echoing the image - return c.encryptResponse("", timestamp, nonce, WeComAIBotStreamResponse{ - MsgType: "stream", - Stream: WeComAIBotStreamInfo{ - ID: c.generateStreamID(), - Finish: true, - Content: fmt.Sprintf( - "Image received (URL: %s), but image messages are not yet supported", - imageURL, - ), - }, - }) -} - -// handleMixedMessage handles mixed (text + image) messages -func (c *WeComAIBotChannel) handleMixedMessage( - ctx context.Context, - msg WeComAIBotMessage, - timestamp, nonce string, -) string { - logger.WarnC("wecom_aibot", "Mixed message type not yet fully implemented") - return c.encryptResponse("", timestamp, nonce, WeComAIBotStreamResponse{ - MsgType: "stream", - Stream: WeComAIBotStreamInfo{ - ID: c.generateStreamID(), - Finish: true, - Content: "Mixed message type is not yet supported", - }, - }) -} - -// handleEventMessage handles event messages -func (c *WeComAIBotChannel) handleEventMessage( - ctx context.Context, - msg WeComAIBotMessage, - timestamp, nonce string, -) string { - eventType := "" - if msg.Event != nil { - eventType = msg.Event.EventType - } - logger.DebugCF("wecom_aibot", "Received event", map[string]any{ - "event_type": eventType, - }) - - // Send welcome message when user opens the chat window - if eventType == "enter_chat" && c.config.WelcomeMessage != "" { - streamID := c.generateStreamID() - return c.encryptResponse(streamID, timestamp, nonce, WeComAIBotStreamResponse{ - MsgType: "stream", - Stream: WeComAIBotStreamInfo{ - ID: streamID, - Finish: true, - Content: c.config.WelcomeMessage, - }, - }) - } - - return c.encryptEmptyResponse(timestamp, nonce) -} - -// getStreamResponse gets the next streaming response for a task. -// - If agent replied: return finish=true with the real answer. -// - If deadline passed: return finish=true with a "please wait" notice, keep task alive for response_url. -// - Otherwise: return finish=false (empty), client will poll again. -func (c *WeComAIBotChannel) getStreamResponse(task *streamTask, timestamp, nonce string) string { - var content string - var finish bool - var closeStreamOnly bool // close stream but do NOT remove task (response_url still pending) - - select { - case answer := <-task.answerCh: - // Agent replied before deadline — normal finish. - content = answer - finish = true - default: - if time.Now().After(task.Deadline) { - // Deadline reached: close the stream with a notice, then wait for agent via response_url. - content = c.config.ProcessingMessage - finish = true - closeStreamOnly = true - logger.InfoCF( - "wecom_aibot", - "Stream deadline reached, switching to response_url mode", - map[string]any{ - "stream_id": task.StreamID, - "chat_id": task.ChatID, - "response_url": task.ResponseURL != "", - }, - ) - } - // else: still waiting, return finish=false - } - - if finish && !closeStreamOnly { - // Normal finish: remove from all maps. - c.removeTask(task) - } else if closeStreamOnly { - // Mark stream as closed and remove from streamTasks under a single lock - // to keep StreamClosed/StreamClosedAt consistent with map membership. - c.taskMu.Lock() - task.StreamClosed = true - task.StreamClosedAt = time.Now() - delete(c.streamTasks, task.StreamID) - c.taskMu.Unlock() - } - - response := WeComAIBotStreamResponse{ - MsgType: "stream", - Stream: WeComAIBotStreamInfo{ - ID: task.StreamID, - Finish: finish, - Content: content, - }, - } - - return c.encryptResponse(task.StreamID, timestamp, nonce, response) -} - -// removeTask removes a task from both streamTasks and chatTasks, marks it finished, -// and cancels its context to interrupt the associated agent goroutine. -func (c *WeComAIBotChannel) removeTask(task *streamTask) { - // Cancel first so the agent goroutine stops as soon as possible, - // before we acquire the write lock. - task.cancel() - - c.taskMu.Lock() - task.Finished = true // written under c.taskMu, consistent with all readers - delete(c.streamTasks, task.StreamID) - queue := c.chatTasks[task.ChatID] - for i, t := range queue { - if t == task { - c.chatTasks[task.ChatID] = append(queue[:i], queue[i+1:]...) - break - } - } - if len(c.chatTasks[task.ChatID]) == 0 { - delete(c.chatTasks, task.ChatID) - } - c.taskMu.Unlock() -} - -// sendViaResponseURL posts a markdown reply to the WeCom response_url. -// response_url is valid for 1 hour and can only be used once per callback. -// Returned errors are wrapped with channels.ErrRateLimit, channels.ErrTemporary, -// or channels.ErrSendFailed so the manager can apply the right retry policy. -func (c *WeComAIBotChannel) sendViaResponseURL(responseURL, content string) error { - payload := map[string]any{ - "msgtype": "markdown", - "markdown": map[string]string{ - "content": content, - }, - } - body, err := json.Marshal(payload) - if err != nil { - return fmt.Errorf("failed to marshal payload: %w", err) - } - - ctx, cancel := context.WithTimeout(c.ctx, 15*time.Second) - defer cancel() - - req, err := http.NewRequestWithContext(ctx, http.MethodPost, responseURL, bytes.NewBuffer(body)) - if err != nil { - return fmt.Errorf("failed to create request: %w", err) - } - req.Header.Set("Content-Type", "application/json; charset=utf-8") - - resp, err := responseURLHTTPClient.Do(req) - if err != nil { - return fmt.Errorf("post to response_url failed: %w: %w", channels.ErrTemporary, err) - } - defer resp.Body.Close() - - if resp.StatusCode == http.StatusOK { - return nil - } - - const maxErrBody = 64 << 10 // 64 KB is more than enough for any error response - respBody, err := io.ReadAll(io.LimitReader(resp.Body, maxErrBody)) - if err != nil { - return fmt.Errorf("reading response_url body: %w: %w", channels.ErrTemporary, err) - } - switch { - case resp.StatusCode == http.StatusTooManyRequests: - return fmt.Errorf("response_url rate limited (%d): %s: %w", - resp.StatusCode, respBody, channels.ErrRateLimit) - case resp.StatusCode >= 500: - return fmt.Errorf("response_url server error (%d): %s: %w", - resp.StatusCode, respBody, channels.ErrTemporary) - default: - return fmt.Errorf("response_url returned %d: %s: %w", - resp.StatusCode, respBody, channels.ErrSendFailed) - } -} - -// encryptResponse encrypts a streaming response -func (c *WeComAIBotChannel) encryptResponse( - streamID, timestamp, nonce string, - response WeComAIBotStreamResponse, -) string { - // Marshal response to JSON - plaintext, err := json.Marshal(response) - if err != nil { - logger.ErrorCF("wecom_aibot", "Failed to marshal response", map[string]any{ - "error": err, - }) - return "" - } - - logger.DebugCF("wecom_aibot", "Encrypting response", map[string]any{ - "stream_id": streamID, - "finish": response.Stream.Finish, - "preview": utils.Truncate(response.Stream.Content, 100), - }) - - // Encrypt message - encrypted, err := c.encryptMessage(string(plaintext), "") - if err != nil { - logger.ErrorCF("wecom_aibot", "Failed to encrypt message", map[string]any{ - "error": err, - }) - return "" - } - - // Generate signature - signature := computeSignature(c.config.Token(), timestamp, nonce, encrypted) - - // Build encrypted response - encryptedResp := WeComAIBotEncryptedResponse{ - Encrypt: encrypted, - MsgSignature: signature, - Timestamp: timestamp, - Nonce: nonce, - } - - respJSON, err := json.Marshal(encryptedResp) - if err != nil { - logger.ErrorCF("wecom_aibot", "Failed to marshal encrypted response", map[string]any{ - "error": err, - }) - return "" - } - - logger.DebugCF("wecom_aibot", "Response encrypted", map[string]any{ - "stream_id": streamID, - }) - - return string(respJSON) -} - -// encryptEmptyResponse returns a minimal valid encrypted response -func (c *WeComAIBotChannel) encryptEmptyResponse(timestamp, nonce string) string { - // Construct a zero-value stream response and encrypt it so that - // WeCom always receives a syntactically valid encrypted JSON object. - emptyResp := WeComAIBotStreamResponse{} - return c.encryptResponse("", timestamp, nonce, emptyResp) -} - -// encryptMessage encrypts a plain text message for WeCom AI Bot -func (c *WeComAIBotChannel) encryptMessage(plaintext, receiveid string) (string, error) { - aesKey, err := decodeWeComAESKey(c.config.EncodingAESKey()) - if err != nil { - return "", err - } - - frame, err := packWeComFrame(plaintext, receiveid) - if err != nil { - return "", err - } - - // PKCS7 padding then AES-CBC encrypt - paddedFrame := pkcs7Pad(frame, blockSize) - ciphertext, err := encryptAESCBC(aesKey, paddedFrame) - if err != nil { - return "", err - } - - return base64.StdEncoding.EncodeToString(ciphertext), nil -} - -// func (c *WeComAIBotChannel) downloadAndDecryptImage( -// ctx context.Context, -// imageURL string, -// ) ([]byte, error) { -// // Download image -// req, err := http.NewRequestWithContext(ctx, http.MethodGet, imageURL, nil) -// if err != nil { -// return nil, fmt.Errorf("failed to create request: %w", err) -// } - -// client := &http.Client{ -// Timeout: 15 * time.Second, -// } - -// resp, err := client.Do(req) -// if err != nil { -// return nil, fmt.Errorf("failed to download image: %w", err) -// } -// defer resp.Body.Close() - -// if resp.StatusCode != http.StatusOK { -// return nil, fmt.Errorf("download failed with status: %d", resp.StatusCode) -// } - -// // Limit image download to 20 MB to prevent memory exhaustion -// const maxImageSize = 20 << 20 // 20 MB -// encryptedData, err := io.ReadAll(io.LimitReader(resp.Body, maxImageSize+1)) -// if err != nil { -// return nil, fmt.Errorf("failed to read image data: %w", err) -// } -// if len(encryptedData) > maxImageSize { -// return nil, fmt.Errorf("image too large (exceeds %d MB)", maxImageSize>>20) -// } - -// logger.DebugCF("wecom_aibot", "Image downloaded", map[string]any{ -// "size": len(encryptedData), -// }) - -// // Decode AES key -// aesKey, err := decodeWeComAESKey(c.config.EncodingAESKey) -// if err != nil { -// return nil, err -// } - -// // Decrypt image (AES-CBC with IV = first 16 bytes of key, PKCS7 padding stripped) -// decryptedData, err := decryptAESCBC(aesKey, encryptedData) -// if err != nil { -// return nil, fmt.Errorf("failed to decrypt image: %w", err) -// } - -// logger.DebugCF("wecom_aibot", "Image decrypted", map[string]any{ -// "size": len(decryptedData), -// }) - -// return decryptedData, nil -// } - -// generateRandomID generates a cryptographically random alphanumeric ID of -// length n. Used for stream IDs and WebSocket request IDs. -func generateRandomID(n int) string { - const letters = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789" - b := make([]byte, n) - for i := range b { - num, _ := rand.Int(rand.Reader, big.NewInt(int64(len(letters)))) - b[i] = letters[num.Int64()] - } - return string(b) -} - -// generateStreamID generates a random 10-character stream ID (webhook mode). -func (c *WeComAIBotChannel) generateStreamID() string { - return generateRandomID(10) -} - -// cleanupLoop periodically cleans up old streaming tasks -func (c *WeComAIBotChannel) cleanupLoop() { - ticker := time.NewTicker(5 * time.Minute) - defer ticker.Stop() - - for { - select { - case <-ticker.C: - c.cleanupOldTasks() - case <-c.ctx.Done(): - return - } - } -} - -// cleanupOldTasks removes tasks that have exceeded their expected lifetime: -// - Active tasks (in streamTasks): cleaned up after 1 hour (response_url validity window). -// - StreamClosed tasks (in chatTasks only): cleaned up after streamClosedGracePeriod. -// These tasks are waiting for the agent to call Send() via response_url. If the agent -// crashes or times out without calling Send(), we must not let them accumulate indefinitely. -// The grace period is generous enough to cover typical LLM latency but far shorter than 1 hour, -// preventing chatTasks from filling up when many requests time out in quick succession. -const ( - streamClosedGracePeriod = 10 * time.Minute // max wait for agent after stream closes - taskMaxLifetime = 1 * time.Hour // absolute max (≈ response_url validity) -) - -func (c *WeComAIBotChannel) cleanupOldTasks() { - c.taskMu.Lock() - defer c.taskMu.Unlock() - - now := time.Now() - cutoff := now.Add(-taskMaxLifetime) - for id, task := range c.streamTasks { - if task.CreatedTime.Before(cutoff) { - delete(c.streamTasks, id) - task.cancel() // interrupt agent goroutine still waiting for LLM - queue := c.chatTasks[task.ChatID] - for i, t := range queue { - if t == task { - c.chatTasks[task.ChatID] = append(queue[:i], queue[i+1:]...) - break - } - } - if len(c.chatTasks[task.ChatID]) == 0 { - delete(c.chatTasks, task.ChatID) - } - logger.DebugCF("wecom_aibot", "Cleaned up expired task", map[string]any{ - "stream_id": id, - }) - } - } - // Clean up StreamClosed tasks from chatTasks. - // Two expiry conditions are checked: - // 1. Absolute expiry: task was created more than taskMaxLifetime ago. - // 2. Grace expiry: stream closed more than streamClosedGracePeriod ago - // (agent had enough time to reply; it is not coming back). - for chatID, queue := range c.chatTasks { - filtered := queue[:0] - for i, t := range queue { - absoluteExpired := t.CreatedTime.Before(cutoff) - graceExpired := t.StreamClosed && - !t.StreamClosedAt.IsZero() && - t.StreamClosedAt.Before(now.Add(-streamClosedGracePeriod)) - if t.Finished { - // Finished tasks should have been removed by removeTask(). - // Finding one here (especially not at position 0) means an - // unexpected code path left it stranded, causing the queue to - // grow silently. Log a warning so it is visible, then drop it. - if i > 0 { - logger.WarnCF("wecom_aibot", - "Found stranded Finished task in the middle of chatTasks queue; "+ - "this should not happen — removeTask() should have spliced it out", - map[string]any{ - "chat_id": chatID, - "stream_id": t.StreamID, - "position": i, - }) - } - // The task is already finished; its context was already canceled - // by removeTask(), so no further action is required. - continue - } else if !absoluteExpired && !graceExpired { - filtered = append(filtered, t) - } else { - t.cancel() // cancel any lingering agent goroutine - } - } - if len(filtered) == 0 { - delete(c.chatTasks, chatID) - } else { - c.chatTasks[chatID] = filtered - } - } -} - -// handleHealth handles health check requests -func (c *WeComAIBotChannel) handleHealth(w http.ResponseWriter, r *http.Request) { - status := "ok" - if !c.IsRunning() { - status = "not running" - } - - w.Header().Set("Content-Type", "application/json") - w.WriteHeader(http.StatusOK) - json.NewEncoder(w).Encode(map[string]string{ - "status": status, - }) -} diff --git a/pkg/channels/wecom/aibot_test.go b/pkg/channels/wecom/aibot_test.go deleted file mode 100644 index 11c4393d6..000000000 --- a/pkg/channels/wecom/aibot_test.go +++ /dev/null @@ -1,559 +0,0 @@ -package wecom - -import ( - "context" - "encoding/json" - "testing" - "time" - - "github.com/sipeed/picoclaw/pkg/bus" - "github.com/sipeed/picoclaw/pkg/channels" - "github.com/sipeed/picoclaw/pkg/config" -) - -// ---- Webhook mode tests ---- - -func TestNewWeComAIBotChannel_WebhookMode(t *testing.T) { - t.Run("success with valid config", func(t *testing.T) { - cfg := config.WeComAIBotConfig{} - cfg.Enabled = true - cfg.SetToken("test_token") - cfg.SetEncodingAESKey("testkey1234567890123456789012345678901234567") - cfg.WebhookPath = "/webhook/test" - - messageBus := bus.NewMessageBus() - ch, err := NewWeComAIBotChannel(cfg, messageBus) - if err != nil { - t.Fatalf("Expected no error, got %v", err) - } - if ch == nil { - t.Fatal("Expected channel to be created") - } - if ch.Name() != "wecom_aibot" { - t.Errorf("Expected name 'wecom_aibot', got '%s'", ch.Name()) - } - // Webhook mode must implement WebhookHandler. - if _, ok := ch.(channels.WebhookHandler); !ok { - t.Error("Webhook mode channel should implement WebhookHandler") - } - }) - - t.Run("error with missing token", func(t *testing.T) { - cfg := config.WeComAIBotConfig{} - cfg.Enabled = true - cfg.SetEncodingAESKey("testkey1234567890123456789012345678901234567") - - messageBus := bus.NewMessageBus() - _, err := NewWeComAIBotChannel(cfg, messageBus) - if err == nil { - t.Fatal("Expected error for missing token, got nil") - } - }) - - t.Run("error with missing encoding key", func(t *testing.T) { - cfg := config.WeComAIBotConfig{} - cfg.Enabled = true - cfg.SetToken("test_token") - - messageBus := bus.NewMessageBus() - _, err := NewWeComAIBotChannel(cfg, messageBus) - if err == nil { - t.Fatal("Expected error for missing encoding key, got nil") - } - }) -} - -func TestWeComAIBotWebhookChannelStartStop(t *testing.T) { - cfg := config.WeComAIBotConfig{ - Enabled: true, - } - cfg.SetToken("test_token") - cfg.SetEncodingAESKey("testkey1234567890123456789012345678901234567") - - messageBus := bus.NewMessageBus() - ch, err := NewWeComAIBotChannel(cfg, messageBus) - if err != nil { - t.Fatalf("Failed to create channel: %v", err) - } - - ctx := context.Background() - - if err := ch.Start(ctx); err != nil { - t.Fatalf("Failed to start channel: %v", err) - } - if !ch.IsRunning() { - t.Error("Expected channel to be running after Start") - } - - if err := ch.Stop(ctx); err != nil { - t.Fatalf("Failed to stop channel: %v", err) - } - if ch.IsRunning() { - t.Error("Expected channel to be stopped after Stop") - } -} - -func TestWeComAIBotChannelWebhookPath(t *testing.T) { - t.Run("default path", func(t *testing.T) { - cfg := config.WeComAIBotConfig{} - cfg.Enabled = true - cfg.SetToken("test_token") - cfg.SetEncodingAESKey("testkey1234567890123456789012345678901234567") - - messageBus := bus.NewMessageBus() - ch, _ := NewWeComAIBotChannel(cfg, messageBus) - - wh, ok := ch.(channels.WebhookHandler) - if !ok { - t.Fatal("Expected channel to implement WebhookHandler") - } - expectedPath := "/webhook/wecom-aibot" - if wh.WebhookPath() != expectedPath { - t.Errorf("Expected webhook path '%s', got '%s'", expectedPath, wh.WebhookPath()) - } - }) - - t.Run("custom path", func(t *testing.T) { - customPath := "/custom/webhook" - cfg := config.WeComAIBotConfig{} - cfg.Enabled = true - cfg.SetToken("test_token") - cfg.SetEncodingAESKey("testkey1234567890123456789012345678901234567") - cfg.WebhookPath = customPath - - messageBus := bus.NewMessageBus() - ch, _ := NewWeComAIBotChannel(cfg, messageBus) - - wh, ok := ch.(channels.WebhookHandler) - if !ok { - t.Fatal("Expected channel to implement WebhookHandler") - } - if wh.WebhookPath() != customPath { - t.Errorf("Expected webhook path '%s', got '%s'", customPath, wh.WebhookPath()) - } - }) -} - -func TestWeComAIBotChannelGetStreamResponseProcessingMessage(t *testing.T) { - validAESKey := "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFG" - - t.Run("uses default processing message", func(t *testing.T) { - cfg := config.WeComAIBotConfig{ - Enabled: true, - } - cfg.SetToken("test_token") - cfg.SetEncodingAESKey(validAESKey) - - messageBus := bus.NewMessageBus() - channel, err := NewWeComAIBotChannel(cfg, messageBus) - if err != nil { - t.Fatalf("Failed to create channel: %v", err) - } - ch, ok := channel.(*WeComAIBotChannel) - if !ok { - t.Fatal("Expected webhook mode channel") - } - - task := &streamTask{ - StreamID: "stream-default", - ChatID: "chat-default", - Deadline: time.Now().Add(-time.Second), - } - ch.streamTasks[task.StreamID] = task - ch.chatTasks[task.ChatID] = []*streamTask{task} - - resp := decodeStreamResponse(t, ch, ch.getStreamResponse(task, "1234567890", "nonce")) - - if !resp.Stream.Finish { - t.Fatal("Expected finished stream response after deadline") - } - if resp.Stream.Content != config.DefaultWeComAIBotProcessingMessage { - t.Fatalf("Expected default processing message %q, got %q", - config.DefaultWeComAIBotProcessingMessage, resp.Stream.Content) - } - if !task.StreamClosed { - t.Fatal("Expected task stream to be marked closed") - } - if _, ok := ch.streamTasks[task.StreamID]; ok { - t.Fatal("Expected closed stream task to be removed from streamTasks") - } - if len(ch.chatTasks[task.ChatID]) != 1 { - t.Fatalf("Expected task to remain queued for response_url delivery, got %d entries", - len(ch.chatTasks[task.ChatID])) - } - }) - - t.Run("uses custom processing message", func(t *testing.T) { - cfg := config.WeComAIBotConfig{ - Enabled: true, - ProcessingMessage: "Please wait a moment. The result will be delivered in a follow-up message.", - } - cfg.SetToken("test_token") - cfg.SetEncodingAESKey(validAESKey) - - messageBus := bus.NewMessageBus() - channel, err := NewWeComAIBotChannel(cfg, messageBus) - if err != nil { - t.Fatalf("Failed to create channel: %v", err) - } - ch, ok := channel.(*WeComAIBotChannel) - if !ok { - t.Fatal("Expected webhook mode channel") - } - - task := &streamTask{ - StreamID: "stream-custom", - ChatID: "chat-custom", - Deadline: time.Now().Add(-time.Second), - } - - resp := decodeStreamResponse(t, ch, ch.getStreamResponse(task, "1234567890", "nonce")) - - if resp.Stream.Content != cfg.ProcessingMessage { - t.Fatalf("Expected custom processing message %q, got %q", cfg.ProcessingMessage, resp.Stream.Content) - } - }) -} - -func TestGenerateStreamID(t *testing.T) { - cfg := config.WeComAIBotConfig{} - cfg.Enabled = true - cfg.SetToken("test_token") - cfg.SetEncodingAESKey("testkey1234567890123456789012345678901234567") - - messageBus := bus.NewMessageBus() - ch, _ := NewWeComAIBotChannel(cfg, messageBus) - webhookCh, ok := ch.(*WeComAIBotChannel) - if !ok { - t.Fatal("Expected webhook mode channel") - } - - ids := make(map[string]bool) - for i := 0; i < 100; i++ { - id := webhookCh.generateStreamID() - if len(id) != 10 { - t.Errorf("Expected stream ID length 10, got %d", len(id)) - } - if ids[id] { - t.Errorf("Duplicate stream ID generated: %s", id) - } - ids[id] = true - } -} - -func TestEncryptDecrypt(t *testing.T) { - // Use a valid 43-character base64 key (企业微信标准格式) - cfg := config.WeComAIBotConfig{} - cfg.Enabled = true - cfg.SetToken("test_token") - cfg.SetEncodingAESKey("abcdefghijklmnopqrstuvwxyz0123456789ABCDEFG") // 43 characters - - messageBus := bus.NewMessageBus() - ch, _ := NewWeComAIBotChannel(cfg, messageBus) - webhookCh, ok := ch.(*WeComAIBotChannel) - if !ok { - t.Fatal("Expected webhook mode channel") - } - - plaintext := "Hello, World!" - receiveid := "" - - encrypted, err := webhookCh.encryptMessage(plaintext, receiveid) - if err != nil { - t.Fatalf("Failed to encrypt message: %v", err) - } - if encrypted == "" { - t.Fatal("Encrypted message is empty") - } - - // Decrypt - decrypted, err := decryptMessageWithVerify(encrypted, cfg.EncodingAESKey(), receiveid) - if err != nil { - t.Fatalf("Failed to decrypt message: %v", err) - } - if decrypted != plaintext { - t.Errorf("Expected decrypted message '%s', got '%s'", plaintext, decrypted) - } -} - -func TestGenerateSignature(t *testing.T) { - token := "test_token" - timestamp := "1234567890" - nonce := "test_nonce" - encrypt := "encrypted_msg" - - signature := computeSignature(token, timestamp, nonce, encrypt) - if signature == "" { - t.Error("Generated signature is empty") - } - if !verifySignature(token, signature, timestamp, nonce, encrypt) { - t.Error("Generated signature does not verify correctly") - } -} - -func decodeStreamResponse(t *testing.T, ch *WeComAIBotChannel, encryptedResponse string) WeComAIBotStreamResponse { - t.Helper() - - var wrapped WeComAIBotEncryptedResponse - if err := json.Unmarshal([]byte(encryptedResponse), &wrapped); err != nil { - t.Fatalf("Failed to unmarshal encrypted response: %v", err) - } - - plaintext, err := decryptMessageWithVerify(wrapped.Encrypt, ch.config.EncodingAESKey(), "") - if err != nil { - t.Fatalf("Failed to decrypt response: %v", err) - } - - var resp WeComAIBotStreamResponse - if err := json.Unmarshal([]byte(plaintext), &resp); err != nil { - t.Fatalf("Failed to unmarshal decrypted response: %v", err) - } - - return resp -} - -// ---- WebSocket long-connection mode tests ---- - -func TestNewWeComAIBotChannel_WSMode(t *testing.T) { - t.Run("success with bot_id and secret", func(t *testing.T) { - cfg := config.WeComAIBotConfig{ - Enabled: true, - BotID: "test_bot_id", - } - cfg.SetSecret("test_secret") - messageBus := bus.NewMessageBus() - ch, err := NewWeComAIBotChannel(cfg, messageBus) - if err != nil { - t.Fatalf("Expected no error, got %v", err) - } - if ch == nil { - t.Fatal("Expected channel to be created") - } - if ch.Name() != "wecom_aibot" { - t.Errorf("Expected name 'wecom_aibot', got '%s'", ch.Name()) - } - // WebSocket mode must NOT implement WebhookHandler. - if _, ok := ch.(channels.WebhookHandler); ok { - t.Error("WebSocket mode channel should NOT implement WebhookHandler") - } - }) - - t.Run("ws mode takes priority over webhook fields", func(t *testing.T) { - cfg := config.WeComAIBotConfig{ - Enabled: true, - BotID: "test_bot_id", - } - cfg.SetSecret("test_secret") - cfg.SetToken("also_set") - cfg.SetEncodingAESKey("testkey1234567890123456789012345678901234567") - messageBus := bus.NewMessageBus() - ch, err := NewWeComAIBotChannel(cfg, messageBus) - if err != nil { - t.Fatalf("Expected no error, got %v", err) - } - if _, ok := ch.(*WeComAIBotWSChannel); !ok { - t.Error("Expected WebSocket mode channel when both BotID+secret and Token+Key are set") - } - }) - - t.Run("error with missing bot_id", func(t *testing.T) { - cfg := config.WeComAIBotConfig{ - Enabled: true, - } - cfg.SetSecret("test_secret") - messageBus := bus.NewMessageBus() - _, err := NewWeComAIBotChannel(cfg, messageBus) - // Missing bot_id alone means neither WS mode nor webhook mode is fully configured. - if err == nil { - t.Fatal("Expected error for missing bot_id, got nil") - } - }) - - t.Run("error with missing secret", func(t *testing.T) { - cfg := config.WeComAIBotConfig{ - Enabled: true, - BotID: "test_bot_id", - } - messageBus := bus.NewMessageBus() - _, err := NewWeComAIBotChannel(cfg, messageBus) - if err == nil { - t.Fatal("Expected error for missing secret, got nil") - } - }) -} - -func TestWeComAIBotWSChannelStartStop(t *testing.T) { - cfg := config.WeComAIBotConfig{ - Enabled: true, - BotID: "test_bot_id", - } - cfg.SetSecret("test_secret") - messageBus := bus.NewMessageBus() - ch, err := NewWeComAIBotChannel(cfg, messageBus) - if err != nil { - t.Fatalf("Failed to create channel: %v", err) - } - - ctx := context.Background() - - // Start launches a background goroutine; it should not block or return an error. - if err := ch.Start(ctx); err != nil { - t.Fatalf("Failed to start channel: %v", err) - } - if !ch.IsRunning() { - t.Error("Expected channel to be running after Start") - } - - // Stop should work regardless of whether the WebSocket actually connected. - if err := ch.Stop(ctx); err != nil { - t.Fatalf("Failed to stop channel: %v", err) - } - if ch.IsRunning() { - t.Error("Expected channel to be stopped after Stop") - } -} - -func TestGenerateRandomID(t *testing.T) { - ids := make(map[string]bool) - for i := 0; i < 200; i++ { - id := generateRandomID(10) - if len(id) != 10 { - t.Errorf("Expected ID length 10, got %d", len(id)) - } - if ids[id] { - t.Errorf("Duplicate ID generated: %s", id) - } - ids[id] = true - } -} - -func TestWSGenerateID(t *testing.T) { - ids := make(map[string]bool) - for i := 0; i < 200; i++ { - id := wsGenerateID() - if len(id) != 10 { - t.Errorf("Expected ID length 10, got %d", len(id)) - } - if ids[id] { - t.Errorf("Duplicate wsGenerateID result: %s", id) - } - ids[id] = true - } -} - -// ---- Webhook streaming fallback tests ---- - -// makeWebhookChannel creates a started WeComAIBotChannel for testing. -func makeWebhookChannel(t *testing.T) *WeComAIBotChannel { - t.Helper() - cfg := config.WeComAIBotConfig{ - Enabled: true, - } - cfg.SetToken("test_token") - cfg.SetEncodingAESKey("abcdefghijklmnopqrstuvwxyz0123456789ABCDEFG") - ch, err := NewWeComAIBotChannel(cfg, bus.NewMessageBus()) - if err != nil { - t.Fatalf("create channel: %v", err) - } - wc := ch.(*WeComAIBotChannel) - wc.ctx, wc.cancel = context.WithCancel(context.Background()) - return wc -} - -// makeStreamTask creates and registers a streamTask for testing. -func makeStreamTask(t *testing.T, ch *WeComAIBotChannel, streamID, chatID string, deadline time.Time) *streamTask { - t.Helper() - task := &streamTask{ - StreamID: streamID, - ChatID: chatID, - Deadline: deadline, - answerCh: make(chan string, 1), - } - task.ctx, task.cancel = context.WithCancel(ch.ctx) - ch.taskMu.Lock() - ch.streamTasks[streamID] = task - ch.chatTasks[chatID] = append(ch.chatTasks[chatID], task) - ch.taskMu.Unlock() - return task -} - -// TestGetStreamResponse_ImmediateAnswer verifies that when the agent has already -// placed its answer in answerCh, getStreamResponse returns a finish=true response -// and fully removes the task. -func TestGetStreamResponse_ImmediateAnswer(t *testing.T) { - ch := makeWebhookChannel(t) - defer ch.cancel() - - task := makeStreamTask(t, ch, "stream-1", "chat-1", time.Now().Add(30*time.Second)) - task.answerCh <- "hello from agent" - - result := ch.getStreamResponse(task, "ts123", "nonce123") - if result == "" { - t.Fatal("expected non-empty encrypted response") - } - - ch.taskMu.RLock() - _, exists := ch.streamTasks["stream-1"] - ch.taskMu.RUnlock() - if exists { - t.Error("task should have been removed from streamTasks after normal finish") - } - if !task.Finished { - t.Error("task.Finished should be true after normal finish") - } -} - -// TestGetStreamResponse_DeadlinePassed verifies that when the stream deadline has -// elapsed (no agent reply yet), getStreamResponse closes the stream but keeps the -// task alive so the response_url fallback can still deliver the answer. -func TestGetStreamResponse_DeadlinePassed(t *testing.T) { - ch := makeWebhookChannel(t) - defer ch.cancel() - - task := makeStreamTask(t, ch, "stream-2", "chat-2", time.Now().Add(-time.Millisecond)) - - result := ch.getStreamResponse(task, "ts456", "nonce456") - if result == "" { - t.Fatal("expected non-empty encrypted response") - } - - ch.taskMu.RLock() - _, stillStreaming := ch.streamTasks["stream-2"] - ch.taskMu.RUnlock() - if stillStreaming { - t.Error("task should have been removed from streamTasks after deadline") - } - if !task.StreamClosed { - t.Error("task.StreamClosed should be true after deadline") - } - if task.Finished { - t.Error("task.Finished must remain false: agent reply still expected via response_url") - } -} - -// TestGetStreamResponse_StillPending verifies that when neither the agent has -// replied nor the deadline has passed, getStreamResponse returns without altering -// task state (client should poll again). -func TestGetStreamResponse_StillPending(t *testing.T) { - ch := makeWebhookChannel(t) - defer ch.cancel() - - task := makeStreamTask(t, ch, "stream-3", "chat-3", time.Now().Add(30*time.Second)) - - result := ch.getStreamResponse(task, "ts789", "nonce789") - if result == "" { - t.Fatal("expected non-empty encrypted response") - } - - ch.taskMu.RLock() - _, exists := ch.streamTasks["stream-3"] - ch.taskMu.RUnlock() - if !exists { - t.Error("pending task should still be in streamTasks") - } - if task.Finished || task.StreamClosed { - t.Error("pending task should not be finished or stream-closed") - } - // Cleanup. - ch.removeTask(task) -} diff --git a/pkg/channels/wecom/aibot_ws.go b/pkg/channels/wecom/aibot_ws.go deleted file mode 100644 index 53dd7071f..000000000 --- a/pkg/channels/wecom/aibot_ws.go +++ /dev/null @@ -1,1347 +0,0 @@ -package wecom - -import ( - "context" - "encoding/base64" - "encoding/json" - "fmt" - "io" - "net/http" - "os" - "path/filepath" - "strings" - "sync" - "time" - - "github.com/gorilla/websocket" - - "github.com/sipeed/picoclaw/pkg/bus" - "github.com/sipeed/picoclaw/pkg/channels" - "github.com/sipeed/picoclaw/pkg/config" - "github.com/sipeed/picoclaw/pkg/identity" - "github.com/sipeed/picoclaw/pkg/logger" - "github.com/sipeed/picoclaw/pkg/media" - "github.com/sipeed/picoclaw/pkg/utils" -) - -// Long-connection WebSocket endpoint. -// Ref: https://developer.work.weixin.qq.com/document/path/101463 -const ( - wsEndpoint = "wss://openws.work.weixin.qq.com" - wsHeartbeatInterval = 30 * time.Second - wsConnectTimeout = 15 * time.Second - wsSubscribeTimeout = 10 * time.Second - wsSendMsgTimeout = 10 * time.Second - wsRespondMsgTimeout = 10 * time.Second - wsWelcomeMsgTimeout = 5 * time.Second // WeCom requires welcome reply within 5 seconds - wsMaxReconnectWait = 60 * time.Second - wsInitialReconnect = time.Second - - // WeCom requires finish=true within 6 minutes of the first stream frame. - // wsStreamTickInterval controls how often we send an in-progress hint. - // wsStreamMaxDuration is a safety margin below the 6-minute hard limit. - wsStreamTickInterval = 30 * time.Second - wsStreamMaxDuration = 5*time.Minute + 30*time.Second - - // wsImageDownloadTimeout caps the time we spend downloading an inbound image. - wsImageDownloadTimeout = 30 * time.Second - - // Keep req_id -> chat route for late fallback pushes after stream window closes. - wsLateReplyRouteTTL = 30 * time.Minute - - // wsStreamMaxContentBytes is the maximum UTF-8 byte length for the content field - // of a single WeCom AI Bot stream / text / markdown frame. - // Ref: https://developer.work.weixin.qq.com/document/path/101463 - wsStreamMaxContentBytes = 20480 -) - -// wsImageHTTPClient is a shared HTTP client for downloading inbound images. -// Reusing it enables connection pooling across multiple image downloads. -var wsImageHTTPClient = &http.Client{Timeout: wsImageDownloadTimeout} - -// WeComAIBotWSChannel implements channels.Channel for WeCom AI Bot using the -// WebSocket long-connection API. -// Unlike the webhook counterpart it does NOT implement WebhookHandler, so the -// HTTP manager will not register any callback URL for it. -type WeComAIBotWSChannel struct { - *channels.BaseChannel - config config.WeComAIBotConfig - ctx context.Context - cancel context.CancelFunc - - // conn is the active WebSocket connection; nil when disconnected. - // All writes are serialized through connMu. - conn *websocket.Conn - connMu sync.Mutex - - // dedupe prevents duplicate message processing (WeCom may re-deliver). - dedupe *MessageDeduplicator - - // reqStates holds per-req_id runtime state. - // It unifies active task state and late-reply fallback routing. - reqStates map[string]*wsReqState - reqStatesMu sync.Mutex - - // reqPending correlates command req_ids with response channels. - // Used only for subscribe/ping command-response pairs. - reqPending map[string]chan wsEnvelope - reqPendingMu sync.Mutex -} - -// wsTask tracks one in-progress agent reply for a single chat turn. -type wsTask struct { - ReqID string // req_id echoed in all replies for this turn - ChatID string - ChatType uint32 - StreamID string // our generated stream.id - answerCh chan string // agent delivers its reply here via Send() - ctx context.Context - cancel context.CancelFunc -} - -type wsReqState struct { - Task *wsTask - Route wsLateReplyRoute -} - -type wsLateReplyRoute struct { - ChatID string - ChatType uint32 - ReadyAt time.Time - ExpiresAt time.Time -} - -// ---- WebSocket protocol types ---- - -// wsEnvelope is the generic JSON envelope for all WebSocket messages. -type wsEnvelope struct { - Cmd string `json:"cmd,omitempty"` - Headers wsHeaders `json:"headers"` - Body json.RawMessage `json:"body,omitempty"` - ErrCode int `json:"errcode,omitempty"` - ErrMsg string `json:"errmsg,omitempty"` -} - -type wsHeaders struct { - ReqID string `json:"req_id"` -} - -// wsCommand is an outgoing request sent over the WebSocket. -type wsCommand struct { - Cmd string `json:"cmd"` - Headers wsHeaders `json:"headers"` - Body any `json:"body,omitempty"` -} - -type wsSendMsgBody struct { - ChatID string `json:"chatid"` - ChatType uint32 `json:"chat_type,omitempty"` - MsgType string `json:"msgtype"` - Markdown *wsMarkdownContent `json:"markdown,omitempty"` -} - -// wsRespondMsgBody is the body for aibot_respond_msg / aibot_respond_welcome_msg. -type wsRespondMsgBody struct { - MsgType string `json:"msgtype"` - Stream *wsStreamContent `json:"stream,omitempty"` - Text *wsTextContent `json:"text,omitempty"` - Markdown *wsMarkdownContent `json:"markdown,omitempty"` - Image *wsImageContent `json:"image,omitempty"` -} - -type wsStreamContent struct { - ID string `json:"id"` - Finish bool `json:"finish"` - Content string `json:"content,omitempty"` -} - -// wsImageContent carries a base64-encoded image payload for outbound messages. -type wsImageContent struct { - Base64 string `json:"base64"` - MD5 string `json:"md5"` -} - -type wsTextContent struct { - Content string `json:"content"` -} - -type wsMarkdownContent struct { - Content string `json:"content"` -} - -// WeComAIBotWSMessage is the decoded body of aibot_msg_callback / -// aibot_event_callback in WebSocket long-connection mode. -// The structure mirrors WeComAIBotMessage but includes extra fields -// that only appear in long-connection callbacks (Voice, AESKey on Image/File). -type WeComAIBotWSMessage struct { - MsgID string `json:"msgid"` - CreateTime int64 `json:"create_time,omitempty"` - AIBotID string `json:"aibotid"` - ChatID string `json:"chatid,omitempty"` - ChatType string `json:"chattype,omitempty"` // "single" | "group" - From struct { - UserID string `json:"userid"` - } `json:"from"` - MsgType string `json:"msgtype"` - Text *struct { - Content string `json:"content"` - } `json:"text,omitempty"` - Image *struct { - URL string `json:"url"` - AESKey string `json:"aeskey,omitempty"` // long-connection: per-resource decrypt key - } `json:"image,omitempty"` - Voice *struct { - Content string `json:"content"` // WeCom transcribes voice to text in callbacks - } `json:"voice,omitempty"` - Mixed *struct { - MsgItem []struct { - MsgType string `json:"msgtype"` - Text *struct { - Content string `json:"content"` - } `json:"text,omitempty"` - Image *struct { - URL string `json:"url"` - AESKey string `json:"aeskey,omitempty"` - } `json:"image,omitempty"` - } `json:"msg_item"` - } `json:"mixed,omitempty"` - Event *struct { - EventType string `json:"eventtype"` - } `json:"event,omitempty"` - File *struct { - URL string `json:"url"` - AESKey string `json:"aeskey,omitempty"` - } `json:"file,omitempty"` - Video *struct { - URL string `json:"url"` - AESKey string `json:"aeskey,omitempty"` - } `json:"video,omitempty"` -} - -// ---- Constructor ---- - -// newWeComAIBotWSChannel creates a WeComAIBotWSChannel for WebSocket mode. -func newWeComAIBotWSChannel( - cfg config.WeComAIBotConfig, - messageBus *bus.MessageBus, -) (*WeComAIBotWSChannel, error) { - if cfg.BotID == "" || cfg.Secret() == "" { - return nil, fmt.Errorf("bot_id and secret are required for WeCom AI Bot WebSocket mode") - } - - base := channels.NewBaseChannel("wecom_aibot", cfg, messageBus, cfg.AllowFrom, - channels.WithReasoningChannelID(cfg.ReasoningChannelID), - ) - - return &WeComAIBotWSChannel{ - BaseChannel: base, - config: cfg, - dedupe: NewMessageDeduplicator(wecomMaxProcessedMessages), - reqStates: make(map[string]*wsReqState), - reqPending: make(map[string]chan wsEnvelope), - }, nil -} - -// ---- Channel interface ---- - -// Name implements channels.Channel. -func (c *WeComAIBotWSChannel) Name() string { return "wecom_aibot" } - -// Start connects to the WeCom WebSocket endpoint and begins message processing. -func (c *WeComAIBotWSChannel) Start(ctx context.Context) error { - logger.InfoC("wecom_aibot", "Starting WeCom AI Bot channel (WebSocket long-connection mode)...") - c.ctx, c.cancel = context.WithCancel(ctx) - c.SetRunning(true) - go c.connectLoop() - logger.InfoC("wecom_aibot", "WeCom AI Bot channel started (WebSocket mode)") - return nil -} - -// Stop shuts down the channel and closes the WebSocket connection. -func (c *WeComAIBotWSChannel) Stop(_ context.Context) error { - logger.InfoC("wecom_aibot", "Stopping WeCom AI Bot channel (WebSocket mode)...") - if c.cancel != nil { - c.cancel() - } - c.connMu.Lock() - if c.conn != nil { - c.conn.Close() - c.conn = nil - } - c.connMu.Unlock() - c.SetRunning(false) - logger.InfoC("wecom_aibot", "WeCom AI Bot channel stopped") - return nil -} - -// Send delivers the agent reply for msg.ChatID. -// The waiting task goroutine picks it up and writes the final stream response. -func (c *WeComAIBotWSChannel) Send(ctx context.Context, msg bus.OutboundMessage) error { - if !c.IsRunning() { - return channels.ErrNotRunning - } - - // msg.ChatID carries the inbound req_id (set by dispatchWSAgentTask). - // For cron-triggered messages, msg.ChatID is the real WeCom chat/user ID - // and there will be no matching entry in reqStates; fall through to proactive push. - task, route, ok := c.getReqState(msg.ChatID) - if !ok { - // No req_id record found — this is a cron/scheduler-originated message. - // Send it as a proactive markdown push using the chat ID directly. - logger.InfoCF("wecom_aibot", "Send: no req_id state, delivering via proactive push (cron/scheduler)", - map[string]any{"chat_id": msg.ChatID}) - if err := c.wsSendActivePush(msg.ChatID, 0, msg.Content); err != nil { - logger.WarnCF("wecom_aibot", "Proactive push failed", - map[string]any{"chat_id": msg.ChatID, "error": err.Error()}) - return fmt.Errorf("websocket delivery failed: %w", channels.ErrSendFailed) - } - return nil - } - - if task == nil { - if time.Now().Before(route.ReadyAt) { - // Keep using aibot_respond_msg within stream window; do not proactively - // push unless wsStreamMaxDuration has elapsed. - logger.WarnCF("wecom_aibot", "Send: stream window still open, skip proactive push", - map[string]any{"req_id": msg.ChatID, "ready_at": route.ReadyAt.Format(time.RFC3339)}) - return nil - } - - if err := c.wsSendActivePush(route.ChatID, route.ChatType, msg.Content); err != nil { - logger.WarnCF("wecom_aibot", "Late reply proactive push failed", - map[string]any{"req_id": msg.ChatID, "chat_id": route.ChatID, "error": err.Error()}) - return fmt.Errorf("websocket delivery failed: %w", channels.ErrSendFailed) - } - logger.InfoCF("wecom_aibot", "Late reply delivered via proactive push", - map[string]any{"req_id": msg.ChatID, "chat_id": route.ChatID, "chat_type": route.ChatType}) - c.deleteReqState(msg.ChatID) - return nil - } - - // Non-blocking fast path: when answerCh has space, deliver without racing - // against task.ctx.Done() (which fires when the task is canceled by a new - // incoming message, but the response must still be sent). - select { - case task.answerCh <- msg.Content: - return nil - default: - } - // answerCh was full; block with cancellation guards. - select { - case task.answerCh <- msg.Content: - case <-task.ctx.Done(): - return nil - case <-ctx.Done(): - return ctx.Err() - } - return nil -} - -// ---- Connection management ---- - -// wsBackoffResetDuration is the minimum duration a WebSocket connection must -// stay up before we reset the reconnect backoff to its initial value. This -// prevents a short burst of failures from causing long waits after later, -// stable connection periods. -const wsBackoffResetDuration = time.Minute - -// connectLoop maintains the WebSocket connection, reconnecting on failure with -// exponential backoff. -func (c *WeComAIBotWSChannel) connectLoop() { - backoff := wsInitialReconnect - for { - select { - case <-c.ctx.Done(): - return - default: - } - - logger.InfoC("wecom_aibot", "Connecting to WeCom WebSocket endpoint...") - start := time.Now() - if err := c.runConnection(); err != nil { - elapsed := time.Since(start) - // If the connection was stable for long enough, reset backoff so that - // a previous burst of failures does not keep us at the maximum delay. - if elapsed >= wsBackoffResetDuration { - backoff = wsInitialReconnect - } - select { - case <-c.ctx.Done(): - return - default: - logger.WarnCF("wecom_aibot", "WebSocket connection lost, reconnecting", - map[string]any{"error": err.Error(), "backoff": backoff.String()}) - select { - case <-time.After(backoff): - case <-c.ctx.Done(): - return - } - if backoff < wsMaxReconnectWait { - backoff *= 2 - if backoff > wsMaxReconnectWait { - backoff = wsMaxReconnectWait - } - } - } - } else { - // Clean exit (context canceled); stop reconnecting. - return - } - } -} - -// runConnection dials, subscribes, and runs the read/heartbeat loops until the -// connection closes or the channel context is canceled. -func (c *WeComAIBotWSChannel) runConnection() error { - dialCtx, dialCancel := context.WithTimeout(c.ctx, wsConnectTimeout) - conn, httpResp, err := websocket.DefaultDialer.DialContext(dialCtx, wsEndpoint, nil) - dialCancel() - if httpResp != nil { - httpResp.Body.Close() - } - if err != nil { - return fmt.Errorf("dial failed: %w", err) - } - - c.connMu.Lock() - c.conn = conn - c.connMu.Unlock() - - defer func() { - c.connMu.Lock() - if c.conn == conn { - c.conn = nil - } - c.connMu.Unlock() - // Cancel any tasks that were started over this connection so their - // agent goroutines do not keep running after the connection is gone. - c.cancelAllTasks() - }() - - // ---- Read loop (must start BEFORE subscribing) ---- - // sendAndWait blocks waiting for the subscribe response on reqPending; - // readLoop is the only goroutine that delivers messages to reqPending. - // Starting readLoop first avoids a deadlock where sendAndWait times out - // because no one reads the server's reply. - readErrCh := make(chan error, 1) - go func() { readErrCh <- c.readLoop(conn) }() - - // ---- Subscribe ---- - reqID := wsGenerateID() - resp, err := c.sendAndWait(conn, reqID, wsCommand{ - Cmd: "aibot_subscribe", - Headers: wsHeaders{ReqID: reqID}, - Body: map[string]string{ - "bot_id": c.config.BotID, - "secret": c.config.Secret(), - }, - }, wsSubscribeTimeout) - if err != nil { - conn.Close() // stop readLoop - <-readErrCh - return fmt.Errorf("subscribe failed: %w", err) - } - if resp.ErrCode != 0 { - conn.Close() - <-readErrCh - return fmt.Errorf("subscribe rejected (errcode=%d): %s", resp.ErrCode, resp.ErrMsg) - } - - logger.InfoC("wecom_aibot", "WebSocket subscription successful") - - // ---- Heartbeat goroutine ---- - hbDone := make(chan struct{}) - go func() { - defer close(hbDone) - c.heartbeatLoop(conn) - }() - - // Wait for the read loop to exit, then tear down the heartbeat. - readErr := <-readErrCh - conn.Close() // signal heartbeat to stop (idempotent) - <-hbDone - return readErr -} - -// sendAndWait registers a pending-response slot, sends cmd, and blocks until -// the matching response arrives or the timeout/context fires. -func (c *WeComAIBotWSChannel) sendAndWait( - conn *websocket.Conn, - reqID string, - cmd wsCommand, - timeout time.Duration, -) (wsEnvelope, error) { - ch := make(chan wsEnvelope, 1) - c.reqPendingMu.Lock() - c.reqPending[reqID] = ch - c.reqPendingMu.Unlock() - - cleanup := func() { - c.reqPendingMu.Lock() - delete(c.reqPending, reqID) - c.reqPendingMu.Unlock() - } - - data, err := json.Marshal(cmd) - if err != nil { - cleanup() - return wsEnvelope{}, fmt.Errorf("marshal command: %w", err) - } - c.connMu.Lock() - err = conn.WriteMessage(websocket.TextMessage, data) - c.connMu.Unlock() - if err != nil { - cleanup() - return wsEnvelope{}, fmt.Errorf("write command: %w", err) - } - - timer := time.NewTimer(timeout) - defer timer.Stop() - select { - case env := <-ch: - return env, nil - case <-timer.C: - cleanup() - return wsEnvelope{}, fmt.Errorf("timeout waiting for response (req_id=%s)", reqID) - case <-c.ctx.Done(): - cleanup() - return wsEnvelope{}, c.ctx.Err() - } -} - -// heartbeatLoop sends a ping every wsHeartbeatInterval until conn is closed. -// It validates the server's pong response via sendAndWait; a failed pong -// triggers a reconnection by closing the connection. -func (c *WeComAIBotWSChannel) heartbeatLoop(conn *websocket.Conn) { - ticker := time.NewTicker(wsHeartbeatInterval) - defer ticker.Stop() - for { - select { - case <-ticker.C: - reqID := wsGenerateID() - resp, err := c.sendAndWait(conn, reqID, wsCommand{ - Cmd: "ping", - Headers: wsHeaders{ReqID: reqID}, - }, wsHeartbeatInterval) - if err != nil { - logger.WarnCF("wecom_aibot", "Heartbeat failed, closing connection", - map[string]any{"error": err.Error()}) - conn.Close() - return - } - if resp.ErrCode != 0 { - logger.WarnCF("wecom_aibot", "Heartbeat rejected", - map[string]any{"errcode": resp.ErrCode, "errmsg": resp.ErrMsg}) - conn.Close() - return - } - logger.DebugCF("wecom_aibot", "Heartbeat pong received", map[string]any{"req_id": reqID}) - case <-c.ctx.Done(): - return - } - } -} - -// readLoop reads WebSocket messages and dispatches them until the connection -// closes or the channel is stopped. -func (c *WeComAIBotWSChannel) readLoop(conn *websocket.Conn) error { - for { - _, raw, err := conn.ReadMessage() - if err != nil { - select { - case <-c.ctx.Done(): - return nil // clean shutdown - default: - return fmt.Errorf("read error: %w", err) - } - } - - var env wsEnvelope - if err := json.Unmarshal(raw, &env); err != nil { - logger.WarnCF("wecom_aibot", "Failed to parse WebSocket message", - map[string]any{"error": err.Error(), "raw": string(raw)}) - continue - } - - // Command responses have an empty Cmd field; forward to any waiting - // sendAndWait() call, or silently drop if no one is waiting (e.g. - // late responses after timeout). - if env.Cmd == "" && env.Headers.ReqID != "" { - c.reqPendingMu.Lock() - ch, ok := c.reqPending[env.Headers.ReqID] - if ok { - delete(c.reqPending, env.Headers.ReqID) - } - c.reqPendingMu.Unlock() - if ok { - ch <- env - } - continue - } - - // Dispatch to appropriate handler in a separate goroutine so the - // read loop is never blocked by a slow agent. - go c.handleEnvelope(env) - } -} - -// ---- Message / event handlers ---- - -// handleEnvelope routes a WebSocket envelope to the right handler. -func (c *WeComAIBotWSChannel) handleEnvelope(env wsEnvelope) { - switch env.Cmd { - case "aibot_msg_callback": - c.handleMsgCallback(env) - case "aibot_event_callback": - c.handleEventCallback(env) - default: - logger.DebugCF("wecom_aibot", "Unhandled WebSocket command", - map[string]any{"cmd": env.Cmd}) - } -} - -// handleMsgCallback processes aibot_msg_callback. -func (c *WeComAIBotWSChannel) handleMsgCallback(env wsEnvelope) { - var msg WeComAIBotWSMessage - if err := json.Unmarshal(env.Body, &msg); err != nil { - logger.WarnCF("wecom_aibot", "Failed to parse msg callback body", - map[string]any{"error": err.Error()}) - return - } - - // Deduplicate by msgid (WeCom may re-deliver on network issues). - if msg.MsgID != "" && !c.dedupe.MarkMessageProcessed(msg.MsgID) { - logger.DebugCF("wecom_aibot", "Duplicate message ignored", - map[string]any{"msgid": msg.MsgID}) - return - } - - reqID := env.Headers.ReqID - switch msg.MsgType { - case "text": - c.handleWSTextMessage(reqID, msg) - case "image": - c.handleWSImageMessage(reqID, msg) - case "voice": - c.handleWSVoiceMessage(reqID, msg) - case "mixed": - c.handleWSMixedMessage(reqID, msg) - case "file": - c.handleWSFileMessage(reqID, msg) - case "video": - c.handleWSVideoMessage(reqID, msg) - default: - logger.WarnCF("wecom_aibot", "Unsupported message type", - map[string]any{"msgtype": msg.MsgType}) - c.wsSendStreamFinish(reqID, wsGenerateID(), - "Unsupported message type: "+msg.MsgType) - } -} - -// handleEventCallback processes aibot_event_callback. -func (c *WeComAIBotWSChannel) handleEventCallback(env wsEnvelope) { - var msg WeComAIBotWSMessage - if err := json.Unmarshal(env.Body, &msg); err != nil { - logger.WarnCF("wecom_aibot", "Failed to parse event callback body", - map[string]any{"error": err.Error()}) - return - } - - // Deduplicate by msgid. - if msg.MsgID != "" && !c.dedupe.MarkMessageProcessed(msg.MsgID) { - logger.DebugCF("wecom_aibot", "Duplicate event ignored", - map[string]any{"msgid": msg.MsgID}) - return - } - - var eventType string - if msg.Event != nil { - eventType = msg.Event.EventType - } - logger.DebugCF("wecom_aibot", "Received event callback", - map[string]any{"event_type": eventType}) - - switch eventType { - case "enter_chat": - if c.config.WelcomeMessage != "" { - c.wsSendWelcomeMsg(env.Headers.ReqID, c.config.WelcomeMessage) - } - case "disconnected_event": - // The server will close this connection after sending this event. - // connectLoop will detect the closure and reconnect automatically. - logger.WarnC("wecom_aibot", - "Received disconnected_event: this connection is being replaced by a newer one") - default: - logger.DebugCF("wecom_aibot", "Unhandled event type", - map[string]any{"event_type": eventType}) - } -} - -// handleWSTextMessage dispatches a plain-text message to the agent and streams -// the reply back over the WebSocket connection. -func (c *WeComAIBotWSChannel) handleWSTextMessage(reqID string, msg WeComAIBotWSMessage) { - if msg.Text == nil { - logger.ErrorC("wecom_aibot", "text message missing text field") - return - } - c.dispatchWSAgentTask(reqID, msg, msg.Text.Content, nil) -} - -// handleWSImageMessage downloads and stores the inbound image, then dispatches -// it to the agent as a media-tagged message. -func (c *WeComAIBotWSChannel) handleWSImageMessage(reqID string, msg WeComAIBotWSMessage) { - if msg.Image == nil { - logger.WarnC("wecom_aibot", "Image message missing image field") - c.wsSendStreamFinish(reqID, wsGenerateID(), "Image message could not be processed.") - return - } - c.wsHandleMediaMessage(reqID, msg, msg.Image.URL, msg.Image.AESKey, "image") -} - -// wsHandleMediaMessage is a shared helper for image, file and video messages. -// It downloads the resource, stores it in MediaStore, and dispatches to the agent. -func (c *WeComAIBotWSChannel) wsHandleMediaMessage( - reqID string, msg WeComAIBotWSMessage, - resourceURL, aesKey, label string, -) { - chatID := wsChatID(msg) - - ctx, cancel := context.WithTimeout(c.ctx, wsImageDownloadTimeout) - defer cancel() - - ref, err := c.storeWSMedia(ctx, chatID, msg.MsgID, resourceURL, aesKey, wsLabelToDefaultExt(label)) - if err != nil { - logger.WarnCF("wecom_aibot", "Failed to download/store WS "+label, - map[string]any{"error": err.Error(), "url": resourceURL}) - c.wsSendStreamFinish(reqID, wsGenerateID(), - strings.ToUpper(label[:1])+label[1:]+" message could not be processed.") - return - } - - c.dispatchWSAgentTask(reqID, msg, "["+label+"]", []string{ref}) -} - -// handleWSMixedMessage handles mixed text+image messages. -// All text parts are collected into the content string; all image parts are -// downloaded and stored in MediaStore before dispatching to the agent. -func (c *WeComAIBotWSChannel) handleWSMixedMessage(reqID string, msg WeComAIBotWSMessage) { - if msg.Mixed == nil { - logger.WarnC("wecom_aibot", "Mixed message has no content") - c.wsSendStreamFinish(reqID, wsGenerateID(), "Mixed message type is not yet fully supported.") - return - } - - chatID := wsChatID(msg) - - ctx, cancel := context.WithTimeout(c.ctx, wsImageDownloadTimeout) - defer cancel() - - var textParts []string - var mediaRefs []string - for _, item := range msg.Mixed.MsgItem { - switch item.MsgType { - case "text": - if item.Text != nil && item.Text.Content != "" { - textParts = append(textParts, item.Text.Content) - } - case "image": - if item.Image != nil { - ref, err := c.storeWSMedia(ctx, chatID, - msg.MsgID+"-"+wsGenerateID(), item.Image.URL, item.Image.AESKey, ".jpg") - if err != nil { - logger.WarnCF("wecom_aibot", "Failed to download/store mixed image", - map[string]any{"error": err.Error()}) - } else { - mediaRefs = append(mediaRefs, ref) - } - } - default: - logger.WarnCF("wecom_aibot", "Unsupported item type in mixed message", - map[string]any{"msgtype": item.MsgType}) - } - } - - if len(textParts) == 0 && len(mediaRefs) == 0 { - logger.WarnC("wecom_aibot", "Mixed message has no usable content") - c.wsSendStreamFinish(reqID, wsGenerateID(), "Mixed message type is not yet fully supported.") - return - } - - content := strings.Join(textParts, "\n") - if content == "" { - content = "[images]" - } - c.dispatchWSAgentTask(reqID, msg, content, mediaRefs) -} - -// dispatchWSAgentTask registers a new agent task, sends the opening stream frame, -// and starts a goroutine that runs the agent and streams the reply back. -// content is the text forwarded to the agent; mediaRefs are optional media -// store references attached to the inbound message. -func (c *WeComAIBotWSChannel) dispatchWSAgentTask( - reqID string, - msg WeComAIBotWSMessage, - content string, - mediaRefs []string, -) { - userID := msg.From.UserID - if userID == "" { - userID = "unknown" - } - // actualChatID is the real WeCom chat/user ID used for peer identification. - // reqID is used as the routing chatID so each turn is independently addressable. - actualChatID := wsChatID(msg) - - streamID := wsGenerateID() - chatType := wsChatTypeValue(msg.ChatType) - taskCtx, taskCancel := context.WithCancel(c.ctx) - - task := &wsTask{ - ReqID: reqID, - ChatID: actualChatID, - ChatType: chatType, - StreamID: streamID, - answerCh: make(chan string, 1), - ctx: taskCtx, - cancel: taskCancel, - } - // Each req_id is unique per WeCom turn; tasks run concurrently, no cancellation. - c.setReqState(reqID, &wsReqState{ - Task: task, - Route: wsLateReplyRoute{ - ChatID: actualChatID, - ChatType: chatType, - ReadyAt: time.Now().Add(wsStreamMaxDuration), - ExpiresAt: time.Now().Add(wsLateReplyRouteTTL), - }, - }) - - logger.DebugCF("wecom_aibot", "Registered new agent task", - map[string]any{"chat_id": actualChatID, "req_id": reqID, "stream_id": streamID}) - - // Send an empty stream opening frame (finish=false) immediately. - c.wsSendStreamChunk(reqID, streamID, false, "") - - go func() { - defer func() { - taskCancel() - c.clearReqTask(reqID, task) - }() - - sender := bus.SenderInfo{ - Platform: "wecom_aibot", - PlatformID: userID, - CanonicalID: identity.BuildCanonicalID("wecom_aibot", userID), - DisplayName: userID, - } - peerKind := "direct" - if msg.ChatType == "group" { - peerKind = "group" - } - peer := bus.Peer{Kind: peerKind, ID: actualChatID} - metadata := map[string]string{ - "channel": "wecom_aibot", - "chat_id": actualChatID, - "chat_type": msg.ChatType, - "msg_type": msg.MsgType, - "msgid": msg.MsgID, - "aibotid": msg.AIBotID, - "stream_id": streamID, - } - // Pass reqID as chatID: OutboundMessage.ChatID = reqID → Send() finds tasks[reqID]. - c.HandleMessage(taskCtx, peer, reqID, userID, reqID, - content, mediaRefs, metadata, sender) - - // Wait for the agent reply. While waiting, send periodic finish=false - // hints so the user knows processing is still in progress. - // WeCom requires finish=true within 6 minutes of the first stream frame; - // wsStreamMaxDuration enforces that limit with a safety margin. - waitHints := []string{ - "⏳ Processing, please wait...", - "⏳ Still processing, please wait...", - "⏳ Almost there, please wait...", - } - ticker := time.NewTicker(wsStreamTickInterval) - defer ticker.Stop() - deadlineTimer := time.NewTimer(wsStreamMaxDuration) - defer deadlineTimer.Stop() - tickCount := 0 - for { - select { - case answer := <-task.answerCh: - // Split the answer into byte-bounded chunks and send as stream frames. - // All but the last carry finish=false; the final frame closes the stream. - chunks := splitWSContent(answer, wsStreamMaxContentBytes) - for i, chunk := range chunks { - c.wsSendStreamChunk(reqID, streamID, i == len(chunks)-1, chunk) - } - c.deleteReqState(reqID) - return - case <-ticker.C: - hint := waitHints[tickCount%len(waitHints)] - tickCount++ - logger.DebugCF("wecom_aibot", "Sending stream progress hint", - map[string]any{"chat_id": actualChatID, "tick": tickCount}) - c.wsSendStreamChunk(reqID, streamID, false, hint) - case <-deadlineTimer.C: - logger.WarnCF("wecom_aibot", - "Stream response deadline reached, closing stream; late reply will be pushed", - map[string]any{"chat_id": actualChatID}) - c.wsSendStreamFinish(reqID, streamID, - "⏳ Processing is taking longer than expected, the response will be sent as a follow-up message.") - return - case <-taskCtx.Done(): - // Give a short grace period so that a response queued in the bus - // just before cancellation can still be delivered. This closes a - // race where a rapid second message cancels this task after the - // agent already published but before Send() wrote to answerCh. - // - // The connection is gone at this point, so we cannot use - // wsSendStreamFinish. Try wsSendActivePush on the (possibly - // already-restored) connection; if that also fails, leave the - // route intact so Send() can push the reply once reconnected. - select { - case answer := <-task.answerCh: - if err := c.wsSendActivePush(task.ChatID, task.ChatType, answer); err != nil { - logger.WarnCF("wecom_aibot", - "Grace-period push failed after task cancellation; reply may be lost", - map[string]any{"req_id": reqID, "chat_id": task.ChatID, "error": err.Error()}) - } else { - c.deleteReqState(reqID) - } - case <-time.After(100 * time.Millisecond): - } - return - } - } - }() -} - -// handleWSVoiceMessage handles voice messages. -// WeCom transcribes voice to text in the callback; if the transcription is -// present it is dispatched as plain text to the agent. -func (c *WeComAIBotWSChannel) handleWSVoiceMessage(reqID string, msg WeComAIBotWSMessage) { - if msg.Voice != nil && msg.Voice.Content != "" { - c.dispatchWSAgentTask(reqID, msg, msg.Voice.Content, nil) - return - } - c.wsSendStreamFinish(reqID, wsGenerateID(), "Voice messages are not yet supported.") -} - -// handleWSFileMessage handles file messages. -func (c *WeComAIBotWSChannel) handleWSFileMessage(reqID string, msg WeComAIBotWSMessage) { - if msg.File == nil { - logger.WarnC("wecom_aibot", "File message missing file field") - c.wsSendStreamFinish(reqID, wsGenerateID(), "File message could not be processed.") - return - } - c.wsHandleMediaMessage(reqID, msg, msg.File.URL, msg.File.AESKey, "file") -} - -// handleWSVideoMessage handles video messages. -func (c *WeComAIBotWSChannel) handleWSVideoMessage(reqID string, msg WeComAIBotWSMessage) { - if msg.Video == nil { - logger.WarnC("wecom_aibot", "Video message missing video field") - c.wsSendStreamFinish(reqID, wsGenerateID(), "Video message could not be processed.") - return - } - c.wsHandleMediaMessage(reqID, msg, msg.Video.URL, msg.Video.AESKey, "video") -} - -// ---- WebSocket write helpers ---- - -// wsSendStreamChunk sends an aibot_respond_msg stream frame. -func (c *WeComAIBotWSChannel) wsSendStreamChunk(reqID, streamID string, finish bool, content string) { - logger.DebugCF("wecom_aibot", "Sending stream chunk", map[string]any{ - "stream_id": streamID, - "finish": finish, - "preview": utils.Truncate(content, 100), - }) - cmd := wsCommand{ - Cmd: "aibot_respond_msg", - Headers: wsHeaders{ReqID: reqID}, - Body: wsRespondMsgBody{ - MsgType: "stream", - Stream: &wsStreamContent{ - ID: streamID, - Finish: finish, - Content: content, - }, - }, - } - if err := c.writeWSAndWait(cmd, wsRespondMsgTimeout); err != nil { - logger.WarnCF("wecom_aibot", "Stream chunk ack failed", map[string]any{ - "req_id": reqID, - "stream_id": streamID, - "finish": finish, - "error": err, - }) - } -} - -// wsSendStreamFinish sends the final aibot_respond_msg frame (finish=true, no images). -func (c *WeComAIBotWSChannel) wsSendStreamFinish(reqID, streamID, content string) { - c.wsSendStreamChunk(reqID, streamID, true, content) -} - -// wsSendWelcomeMsg sends a text welcome message via aibot_respond_welcome_msg. -func (c *WeComAIBotWSChannel) wsSendWelcomeMsg(reqID, content string) { - logger.DebugCF("wecom_aibot", "Sending welcome message", map[string]any{"req_id": reqID}) - cmd := wsCommand{ - Cmd: "aibot_respond_welcome_msg", - Headers: wsHeaders{ReqID: reqID}, - Body: wsRespondMsgBody{ - MsgType: "text", - Text: &wsTextContent{Content: content}, - }, - } - if err := c.writeWSAndWait(cmd, wsWelcomeMsgTimeout); err != nil { - logger.WarnCF("wecom_aibot", "Welcome message ack failed", - map[string]any{"req_id": reqID, "error": err.Error()}) - } -} - -// wsSendActivePush sends a proactive markdown message using aibot_send_msg. -// Long content is automatically split into byte-bounded chunks (≤ wsStreamMaxContentBytes -// each) and delivered as consecutive messages. -// It is used as a fallback for late replies after stream response window expires. -func (c *WeComAIBotWSChannel) wsSendActivePush(chatID string, chatType uint32, content string) error { - if chatID == "" { - return fmt.Errorf("chatid is empty") - } - for _, chunk := range splitWSContent(content, wsStreamMaxContentBytes) { - reqID := wsGenerateID() - if err := c.writeWSAndWait(wsCommand{ - Cmd: "aibot_send_msg", - Headers: wsHeaders{ReqID: reqID}, - Body: wsSendMsgBody{ - ChatID: chatID, - ChatType: chatType, - MsgType: "markdown", - Markdown: &wsMarkdownContent{Content: chunk}, - }, - }, wsSendMsgTimeout); err != nil { - return err - } - } - return nil -} - -// writeWSAndWait writes cmd to the active connection and validates the command response. -func (c *WeComAIBotWSChannel) writeWSAndWait(cmd wsCommand, timeout time.Duration) error { - if cmd.Headers.ReqID == "" { - return fmt.Errorf("req_id is empty") - } - - c.connMu.Lock() - conn := c.conn - c.connMu.Unlock() - if conn == nil { - return fmt.Errorf("websocket not connected") - } - - resp, err := c.sendAndWait(conn, cmd.Headers.ReqID, cmd, timeout) - if err != nil { - return err - } - if resp.ErrCode != 0 { - return fmt.Errorf("%s rejected (errcode=%d): %s", cmd.Cmd, resp.ErrCode, resp.ErrMsg) - } - return nil -} - -// cancelAllTasks cancels every pending agent task; called when the connection drops. -// It also expires each task's stream window (ReadyAt = now) so that when the agent -// eventually delivers its reply via Send(), the message is forwarded via -// wsSendActivePush on the restored connection instead of being silently discarded. -func (c *WeComAIBotWSChannel) cancelAllTasks() { - c.reqStatesMu.Lock() - defer c.reqStatesMu.Unlock() - now := time.Now() - for _, state := range c.reqStates { - if state != nil && state.Task != nil { - state.Task.cancel() - state.Task = nil - // Expire the stream window immediately so Send() uses wsSendActivePush. - state.Route.ReadyAt = now - } - } -} - -func (c *WeComAIBotWSChannel) setReqState(reqID string, state *wsReqState) { - c.reqStatesMu.Lock() - defer c.reqStatesMu.Unlock() - now := time.Now() - for k, v := range c.reqStates { - if v == nil || now.After(v.Route.ExpiresAt) { - delete(c.reqStates, k) - } - } - c.reqStates[reqID] = state -} - -func (c *WeComAIBotWSChannel) getReqState(reqID string) (*wsTask, wsLateReplyRoute, bool) { - c.reqStatesMu.Lock() - defer c.reqStatesMu.Unlock() - state, ok := c.reqStates[reqID] - if !ok || state == nil { - return nil, wsLateReplyRoute{}, false - } - if time.Now().After(state.Route.ExpiresAt) { - delete(c.reqStates, reqID) - return nil, wsLateReplyRoute{}, false - } - return state.Task, state.Route, true -} - -func (c *WeComAIBotWSChannel) deleteReqState(reqID string) { - c.reqStatesMu.Lock() - delete(c.reqStates, reqID) - c.reqStatesMu.Unlock() -} - -func (c *WeComAIBotWSChannel) clearReqTask(reqID string, task *wsTask) { - c.reqStatesMu.Lock() - defer c.reqStatesMu.Unlock() - state, ok := c.reqStates[reqID] - if !ok || state == nil { - return - } - if state.Task == task { - state.Task = nil - } -} - -func wsChatTypeValue(chatType string) uint32 { - if chatType == "group" { - return 2 - } - return 1 -} - -// wsChatID returns the effective chat ID from a WS message. -// For group messages it is msg.ChatID; for single chats it falls back to the sender's UserID. -func wsChatID(msg WeComAIBotWSMessage) string { - if msg.ChatID != "" { - return msg.ChatID - } - return msg.From.UserID -} - -// wsGenerateID generates a random 10-character alphanumeric ID. -// It is package-level (not a method) so it can be shared by both channel modes. -func wsGenerateID() string { - return generateRandomID(10) -} - -// ---- Inbound media download helpers ---- - -// storeWSMedia downloads the resource at resourceURL (with optional AES-CBC -// decryption) and stores it in the MediaStore. The file extension is inferred -// from the HTTP Content-Type response header; defaultExt is used as a fallback -// when the content type is absent or unrecognized. -func (c *WeComAIBotWSChannel) storeWSMedia( - ctx context.Context, - chatID, msgID, resourceURL, aesKey, defaultExt string, -) (string, error) { - store := c.GetMediaStore() - if store == nil { - return "", fmt.Errorf("no media store available") - } - - const maxSize = 20 << 20 // 20 MB - - req, err := http.NewRequestWithContext(ctx, http.MethodGet, resourceURL, nil) - if err != nil { - return "", fmt.Errorf("create request: %w", err) - } - resp, err := wsImageHTTPClient.Do(req) - if err != nil { - return "", fmt.Errorf("download: %w", err) - } - defer resp.Body.Close() - if resp.StatusCode != http.StatusOK { - return "", fmt.Errorf("download HTTP %d", resp.StatusCode) - } - - // Infer file extension from the Content-Type response header. - ext := wsMediaExtFromContentType(resp.Header.Get("Content-Type")) - if ext == "" { - ext = defaultExt - } - - // Buffer the media in memory, bounded to maxSize. - data, err := io.ReadAll(io.LimitReader(resp.Body, int64(maxSize)+1)) - if err != nil { - return "", fmt.Errorf("read media: %w", err) - } - if len(data) > maxSize { - return "", fmt.Errorf("media too large (> %d MB)", maxSize>>20) - } - - // AES-CBC decryption if a key is present. - if aesKey != "" { - key, decErr := base64.StdEncoding.DecodeString(aesKey) - if decErr != nil || len(key) != 32 { - key, decErr = decodeWeComAESKey(aesKey) - if decErr != nil { - return "", fmt.Errorf("decode media AES key: %w", decErr) - } - } - data, err = decryptAESCBC(key, data) - if err != nil { - return "", fmt.Errorf("decrypt media: %w", err) - } - } - - // Write to a temp file. The file is owned by the MediaStore and deleted by - // store.ReleaseAll — no caller-side cleanup needed. - mediaDir := filepath.Join(os.TempDir(), "picoclaw_media") - if err = os.MkdirAll(mediaDir, 0o700); err != nil { - return "", fmt.Errorf("mkdir: %w", err) - } - tmpFile, err := os.CreateTemp(mediaDir, msgID+"-*"+ext) - if err != nil { - return "", fmt.Errorf("create temp file: %w", err) - } - tmpPath := tmpFile.Name() - _, writeErr := tmpFile.Write(data) - closeErr := tmpFile.Close() - if writeErr != nil { - os.Remove(tmpPath) - return "", fmt.Errorf("write media: %w", writeErr) - } - if closeErr != nil { - os.Remove(tmpPath) - return "", fmt.Errorf("close media: %w", closeErr) - } - - scope := channels.BuildMediaScope("wecom_aibot", chatID, msgID) - ref, err := store.Store(tmpPath, media.MediaMeta{ - Filename: msgID + ext, - Source: "wecom_aibot", - CleanupPolicy: media.CleanupPolicyDeleteOnCleanup, - }, scope) - if err != nil { - os.Remove(tmpPath) - return "", fmt.Errorf("store: %w", err) - } - return ref, nil -} - -// wsMediaExtFromContentType returns the lowercase file extension (with leading -// dot) for the given Content-Type value, or "" when the type is unrecognized. -func wsMediaExtFromContentType(contentType string) string { - if contentType == "" { - return "" - } - // Strip parameters (e.g. "image/jpeg; charset=utf-8" → "image/jpeg"). - mt := strings.ToLower(strings.TrimSpace(strings.SplitN(contentType, ";", 2)[0])) - switch mt { - case "image/jpeg", "image/jpg": - return ".jpg" - case "image/png": - return ".png" - case "image/gif": - return ".gif" - case "image/webp": - return ".webp" - case "video/mp4": - return ".mp4" - case "video/mpeg", "video/x-mpeg": - return ".mpeg" - case "video/quicktime": - return ".mov" - case "video/webm": - return ".webm" - case "audio/mpeg", "audio/mp3": - return ".mp3" - case "audio/ogg": - return ".ogg" - case "audio/wav": - return ".wav" - case "application/pdf": - return ".pdf" - case "application/zip": - return ".zip" - case "application/x-rar-compressed", "application/vnd.rar": - return ".rar" - case "text/plain": - return ".txt" - case "application/msword": - return ".doc" - case "application/vnd.openxmlformats-officedocument.wordprocessingml.document": - return ".docx" - case "application/vnd.ms-excel": - return ".xls" - case "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": - return ".xlsx" - case "application/vnd.ms-powerpoint": - return ".ppt" - case "application/vnd.openxmlformats-officedocument.presentationml.presentation": - return ".pptx" - } - return "" -} - -// wsLabelToDefaultExt returns the default file extension for the given media label -// used in wsHandleMediaMessage. It is the fallback when Content-Type detection fails. -func wsLabelToDefaultExt(label string) string { - switch label { - case "image": - return ".jpg" - case "video": - return ".mp4" - default: // "file" and any future labels - return ".bin" - } -} - -// ---- Content length helpers ---- - -// splitWSContent splits content into chunks each fitting within maxBytes UTF-8 -// bytes, preserving code block integrity via channels.SplitMessage. -// When SplitMessage still produces an oversized chunk (e.g. dense CJK content), -// splitAtByteBoundary is applied as a last-resort byte-level fallback. -func splitWSContent(content string, maxBytes int) []string { - if len(content) <= maxBytes { - return []string{content} - } - // SplitMessage works in runes. Use maxBytes as the rune limit: for pure ASCII - // this is exact; for multibyte content the byte verification below catches - // any chunk that still overflows. - chunks := channels.SplitMessage(content, maxBytes) - var result []string - for _, chunk := range chunks { - if len(chunk) <= maxBytes { - result = append(result, chunk) - } else { - // Still too large in bytes (e.g. dense CJK); force-split at UTF-8 boundaries. - result = append(result, splitAtByteBoundary(chunk, maxBytes)...) - } - } - return result -} - -// splitAtByteBoundary splits s into parts each ≤ maxBytes bytes by walking back -// from the hard byte limit to find a valid UTF-8 rune start boundary. -// This is a last-resort fallback; it does not try to preserve code blocks. -func splitAtByteBoundary(s string, maxBytes int) []string { - var parts []string - for len(s) > maxBytes { - end := maxBytes - // Walk back past any UTF-8 continuation bytes (high two bits == 10). - for end > 0 && s[end]>>6 == 0b10 { - end-- - } - if end == 0 { - end = maxBytes // shouldn't happen with valid UTF-8 - } - parts = append(parts, s[:end]) - s = strings.TrimLeft(s[end:], " \t\n\r") - } - if s != "" { - parts = append(parts, s) - } - return parts -} diff --git a/pkg/channels/wecom/aibot_ws_test.go b/pkg/channels/wecom/aibot_ws_test.go deleted file mode 100644 index f2f8833a1..000000000 --- a/pkg/channels/wecom/aibot_ws_test.go +++ /dev/null @@ -1,295 +0,0 @@ -package wecom - -import ( - "bytes" - "context" - "net/http" - "net/http/httptest" - "os" - "strings" - "testing" - - "github.com/sipeed/picoclaw/pkg/bus" - "github.com/sipeed/picoclaw/pkg/channels" - "github.com/sipeed/picoclaw/pkg/config" - "github.com/sipeed/picoclaw/pkg/media" -) - -// newTestWSChannel creates a WeComAIBotWSChannel ready for unit testing. -func newTestWSChannel(t *testing.T) *WeComAIBotWSChannel { - t.Helper() - cfg := config.WeComAIBotConfig{ - Enabled: true, - BotID: "test_bot_id", - } - cfg.SetSecret("test_secret") - ch, err := newWeComAIBotWSChannel(cfg, bus.NewMessageBus()) - if err != nil { - t.Fatalf("create WS channel: %v", err) - } - return ch -} - -// TestStoreWSMedia_NilStore verifies that storeWSMedia returns an error when no -// MediaStore has been injected. -func TestStoreWSMedia_NilStore(t *testing.T) { - ch := newTestWSChannel(t) - _, err := ch.storeWSMedia(context.Background(), "chat1", "msg1", "http://any", "", ".jpg") - if err == nil { - t.Fatal("expected error when no MediaStore is set") - } -} - -// TestStoreWSMedia_HTTPError verifies that storeWSMedia propagates HTTP errors -// from the media server. -func TestStoreWSMedia_HTTPError(t *testing.T) { - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - http.Error(w, "not found", http.StatusNotFound) - })) - defer srv.Close() - - ch := newTestWSChannel(t) - ch.SetMediaStore(media.NewFileMediaStore()) - - _, err := ch.storeWSMedia(context.Background(), "chat1", "msg1", srv.URL, "", ".jpg") - if err == nil { - t.Fatal("expected error for HTTP 404") - } -} - -// TestStoreWSMedia_ServerUnavailable verifies that storeWSMedia returns a clear -// error when the media server cannot be reached. -func TestStoreWSMedia_ServerUnavailable(t *testing.T) { - ch := newTestWSChannel(t) - ch.SetMediaStore(media.NewFileMediaStore()) - - // Port 1 is reserved and will refuse the connection immediately. - _, err := ch.storeWSMedia(context.Background(), "chat1", "msg1", "http://127.0.0.1:1", "", ".jpg") - if err == nil { - t.Fatal("expected error for unreachable server") - } -} - -// TestStoreWSMedia_Success_NoAES verifies the happy path: the media is downloaded, -// a media ref is returned, and the file persists and is readable via Resolve until -// ReleaseAll is called. The server returns no Content-Type, so the defaultExt is used. -func TestStoreWSMedia_Success_NoAES(t *testing.T) { - imageData := bytes.Repeat([]byte("x"), 256) - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - w.WriteHeader(http.StatusOK) - _, _ = w.Write(imageData) - })) - defer srv.Close() - - ch := newTestWSChannel(t) - store := media.NewFileMediaStore() - ch.SetMediaStore(store) - - ref, err := ch.storeWSMedia(context.Background(), "chat1", "msg1", srv.URL, "", ".jpg") - if err != nil { - t.Fatalf("expected no error, got %v", err) - } - if ref == "" { - t.Fatal("expected non-empty ref") - } - - // File must be accessible after storeWSMedia returns (no premature deletion). - path, err := store.Resolve(ref) - if err != nil { - t.Fatalf("ref should resolve: %v", err) - } - got, err := os.ReadFile(path) - if err != nil { - t.Fatalf("file should exist at %s: %v", path, err) - } - if !bytes.Equal(got, imageData) { - t.Errorf("content mismatch: got len=%d, want len=%d", len(got), len(imageData)) - } - - // ReleaseAll must delete the file (store owns lifecycle). - scope := channels.BuildMediaScope("wecom_aibot", "chat1", "msg1") - if err := store.ReleaseAll(scope); err != nil { - t.Fatalf("ReleaseAll failed: %v", err) - } - if _, err := os.Stat(path); !os.IsNotExist(err) { - t.Errorf("file should have been deleted by ReleaseAll, stat err: %v", err) - } -} - -// TestStoreWSMedia_MultipleMessages verifies that concurrent media messages with -// different msgIDs do not collide and each resolve to distinct files. -func TestStoreWSMedia_MultipleMessages(t *testing.T) { - imageA := bytes.Repeat([]byte("a"), 64) - imageB := bytes.Repeat([]byte("b"), 64) - - srvA := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - w.WriteHeader(http.StatusOK) - _, _ = w.Write(imageA) - })) - defer srvA.Close() - srvB := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - w.WriteHeader(http.StatusOK) - _, _ = w.Write(imageB) - })) - defer srvB.Close() - - ch := newTestWSChannel(t) - store := media.NewFileMediaStore() - ch.SetMediaStore(store) - - refA, err := ch.storeWSMedia(context.Background(), "chat1", "msgA", srvA.URL, "", ".jpg") - if err != nil { - t.Fatalf("storeWSMedia A: %v", err) - } - refB, err := ch.storeWSMedia(context.Background(), "chat1", "msgB", srvB.URL, "", ".jpg") - if err != nil { - t.Fatalf("storeWSMedia B: %v", err) - } - if refA == refB { - t.Fatal("distinct messages must produce distinct refs") - } - - pathA, _ := store.Resolve(refA) - pathB, _ := store.Resolve(refB) - if pathA == pathB { - t.Fatal("distinct messages must be stored at distinct paths") - } - - gotA, _ := os.ReadFile(pathA) - gotB, _ := os.ReadFile(pathB) - if !bytes.Equal(gotA, imageA) { - t.Errorf("content mismatch for message A") - } - if !bytes.Equal(gotB, imageB) { - t.Errorf("content mismatch for message B") - } -} - -// TestStoreWSMedia_ContentTypeExt verifies that the file extension is inferred -// from the HTTP Content-Type header and the defaultExt fallback is used when the -// type is absent or unrecognized. -func TestStoreWSMedia_ContentTypeExt(t *testing.T) { - tests := []struct { - contentType string - wantExt string - }{ - {"image/jpeg", ".jpg"}, - {"image/png", ".png"}, - {"video/mp4", ".mp4"}, - {"application/pdf", ".pdf"}, - {"application/zip", ".zip"}, - // With parameters stripped. - {"video/mp4; codecs=avc1", ".mp4"}, - // Unknown type → falls back to defaultExt. - {"", ""}, - {"application/octet-stream", ""}, - } - for _, tc := range tests { - got := wsMediaExtFromContentType(tc.contentType) - if got != tc.wantExt { - t.Errorf("wsMediaExtFromContentType(%q) = %q, want %q", tc.contentType, got, tc.wantExt) - } - } - - // End-to-end: server returns Content-Type: video/mp4, defaultExt is .bin. - // The stored file should carry the .mp4 extension, not .bin. - payload := bytes.Repeat([]byte("v"), 128) - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - w.Header().Set("Content-Type", "video/mp4") - w.WriteHeader(http.StatusOK) - _, _ = w.Write(payload) - })) - defer srv.Close() - - ch := newTestWSChannel(t) - store := media.NewFileMediaStore() - ch.SetMediaStore(store) - - ref, err := ch.storeWSMedia(context.Background(), "chat1", "vid1", srv.URL, "", ".bin") - if err != nil { - t.Fatalf("storeWSMedia: %v", err) - } - path, err := store.Resolve(ref) - if err != nil { - t.Fatalf("resolve: %v", err) - } - if ext := path[len(path)-4:]; ext != ".mp4" { - t.Errorf("expected .mp4 extension from Content-Type, got %q", ext) - } -} - -// TestSplitWSContent verifies byte-aware splitting of stream content. -func TestSplitWSContent(t *testing.T) { - t.Run("short content is not split", func(t *testing.T) { - chunks := splitWSContent("hello", 20480) - if len(chunks) != 1 || chunks[0] != "hello" { - t.Fatalf("unexpected chunks: %v", chunks) - } - }) - - t.Run("ASCII content split at byte boundary", func(t *testing.T) { - // Build a string just over the limit. - content := strings.Repeat("a", 20481) - chunks := splitWSContent(content, 20480) - if len(chunks) < 2 { - t.Fatalf("expected >= 2 chunks, got %d", len(chunks)) - } - for i, c := range chunks { - if len(c) > 20480 { - t.Errorf("chunk %d has %d bytes, want <= 20480", i, len(c)) - } - } - // Reassembled content must equal the original (possibly without leading - // whitespace that splitWSContent trims between chunks). - joined := strings.Join(chunks, "") - if len(joined) < len(content)-len(chunks) { - t.Errorf("joined length %d too short (original %d)", len(joined), len(content)) - } - }) - - t.Run("CJK content split within byte limit", func(t *testing.T) { - // Each CJK rune is 3 bytes in UTF-8. - // 7000 CJK chars = 21000 bytes, which exceeds 20480. - content := strings.Repeat("\u4e2d", 7000) - chunks := splitWSContent(content, 20480) - if len(chunks) < 2 { - t.Fatalf("expected >= 2 chunks for 21000-byte CJK content, got %d", len(chunks)) - } - for i, c := range chunks { - if len(c) > 20480 { - t.Errorf("chunk %d has %d bytes, want <= 20480", i, len(c)) - } - // Every chunk must be valid UTF-8. - if !strings.ContainsRune(c, '\u4e2d') && len(c) > 0 { - // quick plausibility check — content was pure CJK - } - } - }) -} - -// TestSplitAtByteBoundary verifies the last-resort byte-boundary splitter. -func TestSplitAtByteBoundary(t *testing.T) { - t.Run("ASCII fits in one chunk", func(t *testing.T) { - parts := splitAtByteBoundary("hello world", 100) - if len(parts) != 1 { - t.Fatalf("expected 1 part, got %d", len(parts)) - } - }) - - t.Run("splits at byte boundary, never mid-rune", func(t *testing.T) { - // 10 CJK characters = 30 bytes; split at 20 bytes. - s := strings.Repeat("\u6587", 10) // 10 × 3 bytes = 30 bytes - parts := splitAtByteBoundary(s, 20) - for i, p := range parts { - if len(p) > 20 { - t.Errorf("part %d has %d bytes, want <= 20", i, len(p)) - } - // Must be valid UTF-8 (no torn multi-byte sequences). - for j, r := range p { - if r == '\uFFFD' { - t.Errorf("part %d has replacement rune at position %d: torn UTF-8", i, j) - } - } - } - }) -} diff --git a/pkg/channels/wecom/app.go b/pkg/channels/wecom/app.go deleted file mode 100644 index fccfc60a3..000000000 --- a/pkg/channels/wecom/app.go +++ /dev/null @@ -1,756 +0,0 @@ -package wecom - -import ( - "bytes" - "context" - "encoding/json" - "encoding/xml" - "fmt" - "io" - "mime/multipart" - "net/http" - "net/url" - "os" - "path/filepath" - "strings" - "sync" - "time" - - "github.com/sipeed/picoclaw/pkg/bus" - "github.com/sipeed/picoclaw/pkg/channels" - "github.com/sipeed/picoclaw/pkg/config" - "github.com/sipeed/picoclaw/pkg/identity" - "github.com/sipeed/picoclaw/pkg/logger" - "github.com/sipeed/picoclaw/pkg/utils" -) - -const ( - wecomAPIBase = "https://qyapi.weixin.qq.com" -) - -// WeComAppChannel implements the Channel interface for WeCom App (企业微信自建应用) -type WeComAppChannel struct { - *channels.BaseChannel - config config.WeComAppConfig - client *http.Client - accessToken string - tokenExpiry time.Time - tokenMu sync.RWMutex - ctx context.Context - cancel context.CancelFunc - processedMsgs *MessageDeduplicator -} - -// WeComXMLMessage represents the XML message structure from WeCom -type WeComXMLMessage struct { - XMLName xml.Name `xml:"xml"` - ToUserName string `xml:"ToUserName"` - FromUserName string `xml:"FromUserName"` - CreateTime int64 `xml:"CreateTime"` - MsgType string `xml:"MsgType"` - Content string `xml:"Content"` - MsgId int64 `xml:"MsgId"` - AgentID int64 `xml:"AgentID"` - PicUrl string `xml:"PicUrl"` - MediaId string `xml:"MediaId"` - Format string `xml:"Format"` - ThumbMediaId string `xml:"ThumbMediaId"` - LocationX float64 `xml:"Location_X"` - LocationY float64 `xml:"Location_Y"` - Scale int `xml:"Scale"` - Label string `xml:"Label"` - Title string `xml:"Title"` - Description string `xml:"Description"` - Url string `xml:"Url"` - Event string `xml:"Event"` - EventKey string `xml:"EventKey"` -} - -// WeComTextMessage represents text message for sending -type WeComTextMessage struct { - ToUser string `json:"touser"` - MsgType string `json:"msgtype"` - AgentID int64 `json:"agentid"` - Text struct { - Content string `json:"content"` - } `json:"text"` - Safe int `json:"safe,omitempty"` -} - -// WeComMarkdownMessage represents markdown message for sending -type WeComMarkdownMessage struct { - ToUser string `json:"touser"` - MsgType string `json:"msgtype"` - AgentID int64 `json:"agentid"` - Markdown struct { - Content string `json:"content"` - } `json:"markdown"` -} - -// WeComImageMessage represents image message for sending -type WeComImageMessage struct { - ToUser string `json:"touser"` - MsgType string `json:"msgtype"` - AgentID int64 `json:"agentid"` - Image struct { - MediaID string `json:"media_id"` - } `json:"image"` -} - -// WeComAccessTokenResponse represents the access token API response -type WeComAccessTokenResponse struct { - ErrCode int `json:"errcode"` - ErrMsg string `json:"errmsg"` - AccessToken string `json:"access_token"` - ExpiresIn int `json:"expires_in"` -} - -// WeComSendMessageResponse represents the send message API response -type WeComSendMessageResponse struct { - ErrCode int `json:"errcode"` - ErrMsg string `json:"errmsg"` - InvalidUser string `json:"invaliduser"` - InvalidParty string `json:"invalidparty"` - InvalidTag string `json:"invalidtag"` -} - -// PKCS7Padding adds PKCS7 padding -type PKCS7Padding struct{} - -// NewWeComAppChannel creates a new WeCom App channel instance -func NewWeComAppChannel(cfg config.WeComAppConfig, messageBus *bus.MessageBus) (*WeComAppChannel, error) { - if cfg.CorpID == "" || cfg.CorpSecret() == "" || cfg.AgentID == 0 { - return nil, fmt.Errorf("wecom_app corp_id, corp_secret and agent_id are required") - } - - base := channels.NewBaseChannel("wecom_app", cfg, messageBus, cfg.AllowFrom, - channels.WithMaxMessageLength(2048), - channels.WithGroupTrigger(cfg.GroupTrigger), - channels.WithReasoningChannelID(cfg.ReasoningChannelID), - ) - - // Client timeout must be >= the configured ReplyTimeout so the - // per-request context deadline is always the effective limit. - clientTimeout := 30 * time.Second - if d := time.Duration(cfg.ReplyTimeout) * time.Second; d > clientTimeout { - clientTimeout = d - } - - ctx, cancel := context.WithCancel(context.Background()) - return &WeComAppChannel{ - BaseChannel: base, - config: cfg, - client: &http.Client{Timeout: clientTimeout}, - ctx: ctx, - cancel: cancel, - processedMsgs: NewMessageDeduplicator(wecomMaxProcessedMessages), - }, nil -} - -// Name returns the channel name -func (c *WeComAppChannel) Name() string { - return "wecom_app" -} - -// Start initializes the WeCom App channel -func (c *WeComAppChannel) Start(ctx context.Context) error { - logger.InfoC("wecom_app", "Starting WeCom App channel...") - - // Cancel the context created in the constructor to avoid a resource leak. - if c.cancel != nil { - c.cancel() - } - c.ctx, c.cancel = context.WithCancel(ctx) - - // Get initial access token - if err := c.refreshAccessToken(); err != nil { - logger.WarnCF("wecom_app", "Failed to get initial access token", map[string]any{ - "error": err.Error(), - }) - } - - // Start token refresh goroutine - go c.tokenRefreshLoop() - - c.SetRunning(true) - logger.InfoC("wecom_app", "WeCom App channel started") - - return nil -} - -// Stop gracefully stops the WeCom App channel -func (c *WeComAppChannel) Stop(ctx context.Context) error { - logger.InfoC("wecom_app", "Stopping WeCom App channel...") - - if c.cancel != nil { - c.cancel() - } - - c.SetRunning(false) - logger.InfoC("wecom_app", "WeCom App channel stopped") - return nil -} - -// Send sends a message to WeCom user proactively using access token -func (c *WeComAppChannel) Send(ctx context.Context, msg bus.OutboundMessage) error { - if !c.IsRunning() { - return channels.ErrNotRunning - } - - accessToken := c.getAccessToken() - if accessToken == "" { - return fmt.Errorf("no valid access token available") - } - - logger.DebugCF("wecom_app", "Sending message", map[string]any{ - "chat_id": msg.ChatID, - "preview": utils.Truncate(msg.Content, 100), - }) - - return c.sendTextMessage(ctx, accessToken, msg.ChatID, msg.Content) -} - -// SendMedia implements the channels.MediaSender interface. -func (c *WeComAppChannel) SendMedia(ctx context.Context, msg bus.OutboundMediaMessage) error { - if !c.IsRunning() { - return channels.ErrNotRunning - } - - accessToken := c.getAccessToken() - if accessToken == "" { - return fmt.Errorf("no valid access token available: %w", channels.ErrTemporary) - } - - store := c.GetMediaStore() - if store == nil { - return fmt.Errorf("no media store available: %w", channels.ErrSendFailed) - } - - for _, part := range msg.Parts { - localPath, err := store.Resolve(part.Ref) - if err != nil { - logger.ErrorCF("wecom_app", "Failed to resolve media ref", map[string]any{ - "ref": part.Ref, - "error": err.Error(), - }) - continue - } - - // Map part type to WeCom media type - var mediaType string - switch part.Type { - case "image": - mediaType = "image" - case "audio": - mediaType = "voice" - case "video": - mediaType = "video" - default: - mediaType = "file" - } - - // Upload media to get media_id - mediaID, err := c.uploadMedia(ctx, accessToken, mediaType, localPath) - if err != nil { - logger.ErrorCF("wecom_app", "Failed to upload media", map[string]any{ - "type": mediaType, - "error": err.Error(), - }) - // Fallback: send caption as text - if part.Caption != "" { - _ = c.sendTextMessage(ctx, accessToken, msg.ChatID, part.Caption) - } - continue - } - - // Send media message using the media_id - if mediaType == "image" { - err = c.sendImageMessage(ctx, accessToken, msg.ChatID, mediaID) - } else { - // For non-image types, send as text fallback with caption - caption := part.Caption - if caption == "" { - caption = fmt.Sprintf("[%s: %s]", part.Type, part.Filename) - } - err = c.sendTextMessage(ctx, accessToken, msg.ChatID, caption) - } - - if err != nil { - return err - } - } - - return nil -} - -// uploadMedia uploads a local file to WeCom temporary media storage. -func (c *WeComAppChannel) uploadMedia(ctx context.Context, accessToken, mediaType, localPath string) (string, error) { - apiURL := fmt.Sprintf("%s/cgi-bin/media/upload?access_token=%s&type=%s", - wecomAPIBase, url.QueryEscape(accessToken), url.QueryEscape(mediaType)) - - file, err := os.Open(localPath) - if err != nil { - return "", fmt.Errorf("failed to open file: %w", err) - } - defer file.Close() - - body := &bytes.Buffer{} - writer := multipart.NewWriter(body) - - filename := filepath.Base(localPath) - formFile, err := writer.CreateFormFile("media", filename) - if err != nil { - return "", fmt.Errorf("failed to create form file: %w", err) - } - - if _, err = io.Copy(formFile, file); err != nil { - return "", fmt.Errorf("failed to copy file content: %w", err) - } - writer.Close() - - req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, body) - if err != nil { - return "", fmt.Errorf("failed to create request: %w", err) - } - req.Header.Set("Content-Type", writer.FormDataContentType()) - - resp, err := c.client.Do(req) - if err != nil { - return "", channels.ClassifyNetError(err) - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - respBody, readErr := io.ReadAll(resp.Body) - if readErr != nil { - return "", channels.ClassifySendError( - resp.StatusCode, - fmt.Errorf("reading wecom upload error response: %w", readErr), - ) - } - return "", channels.ClassifySendError( - resp.StatusCode, - fmt.Errorf("wecom upload error: %s", string(respBody)), - ) - } - - var result struct { - ErrCode int `json:"errcode"` - ErrMsg string `json:"errmsg"` - MediaID string `json:"media_id"` - } - if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { - return "", fmt.Errorf("failed to parse upload response: %w", err) - } - - if result.ErrCode != 0 { - return "", fmt.Errorf("upload API error: %s (code: %d)", result.ErrMsg, result.ErrCode) - } - - return result.MediaID, nil -} - -// sendWeComMessage marshals payload and POSTs it to the WeCom message API. -func (c *WeComAppChannel) sendWeComMessage(ctx context.Context, accessToken string, payload any) error { - apiURL := fmt.Sprintf("%s/cgi-bin/message/send?access_token=%s", wecomAPIBase, accessToken) - - jsonData, err := json.Marshal(payload) - if err != nil { - return fmt.Errorf("failed to marshal message: %w", err) - } - - timeout := c.config.ReplyTimeout - if timeout <= 0 { - timeout = 5 - } - - reqCtx, cancel := context.WithTimeout(ctx, time.Duration(timeout)*time.Second) - defer cancel() - - req, err := http.NewRequestWithContext(reqCtx, http.MethodPost, apiURL, bytes.NewBuffer(jsonData)) - if err != nil { - return fmt.Errorf("failed to create request: %w", err) - } - req.Header.Set("Content-Type", "application/json") - - resp, err := c.client.Do(req) - if err != nil { - return channels.ClassifyNetError(err) - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - respBody, readErr := io.ReadAll(resp.Body) - if readErr != nil { - return channels.ClassifySendError( - resp.StatusCode, - fmt.Errorf("reading wecom_app error response: %w", readErr), - ) - } - return channels.ClassifySendError( - resp.StatusCode, - fmt.Errorf("wecom_app API error: %s", string(respBody)), - ) - } - - respBody, err := io.ReadAll(resp.Body) - if err != nil { - return fmt.Errorf("failed to read response: %w", err) - } - - var sendResp WeComSendMessageResponse - if err := json.Unmarshal(respBody, &sendResp); err != nil { - return fmt.Errorf("failed to parse response: %w", err) - } - - if sendResp.ErrCode != 0 { - return fmt.Errorf("API error: %s (code: %d)", sendResp.ErrMsg, sendResp.ErrCode) - } - - return nil -} - -// sendImageMessage sends an image message using a media_id. -func (c *WeComAppChannel) sendImageMessage(ctx context.Context, accessToken, userID, mediaID string) error { - msg := WeComImageMessage{ - ToUser: userID, - MsgType: "image", - AgentID: c.config.AgentID, - } - msg.Image.MediaID = mediaID - return c.sendWeComMessage(ctx, accessToken, msg) -} - -// WebhookPath returns the path for registering on the shared HTTP server. -func (c *WeComAppChannel) WebhookPath() string { - if c.config.WebhookPath != "" { - return c.config.WebhookPath - } - return "/webhook/wecom-app" -} - -// ServeHTTP implements http.Handler for the shared HTTP server. -func (c *WeComAppChannel) ServeHTTP(w http.ResponseWriter, r *http.Request) { - c.handleWebhook(w, r) -} - -// HealthPath returns the health check endpoint path. -func (c *WeComAppChannel) HealthPath() string { - return "/health/wecom-app" -} - -// HealthHandler handles health check requests. -func (c *WeComAppChannel) HealthHandler(w http.ResponseWriter, r *http.Request) { - c.handleHealth(w, r) -} - -// handleWebhook handles incoming webhook requests from WeCom -func (c *WeComAppChannel) handleWebhook(w http.ResponseWriter, r *http.Request) { - ctx := r.Context() - - // Log all incoming requests for debugging - logger.DebugCF("wecom_app", "Received webhook request", map[string]any{ - "method": r.Method, - "url": r.URL.String(), - "path": r.URL.Path, - "query": r.URL.RawQuery, - }) - - if r.Method == http.MethodGet { - // Handle verification request - c.handleVerification(ctx, w, r) - return - } - - if r.Method == http.MethodPost { - // Handle message callback - c.handleMessageCallback(ctx, w, r) - return - } - - logger.WarnCF("wecom_app", "Method not allowed", map[string]any{ - "method": r.Method, - }) - http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) -} - -// handleVerification handles the URL verification request from WeCom -func (c *WeComAppChannel) handleVerification(ctx context.Context, w http.ResponseWriter, r *http.Request) { - query := r.URL.Query() - msgSignature := query.Get("msg_signature") - timestamp := query.Get("timestamp") - nonce := query.Get("nonce") - echostr := query.Get("echostr") - - logger.DebugCF("wecom_app", "Handling verification request", map[string]any{ - "msg_signature": msgSignature, - "timestamp": timestamp, - "nonce": nonce, - "echostr": echostr, - "corp_id": c.config.CorpID, - }) - - if msgSignature == "" || timestamp == "" || nonce == "" || echostr == "" { - logger.ErrorC("wecom_app", "Missing parameters in verification request") - http.Error(w, "Missing parameters", http.StatusBadRequest) - return - } - - // Verify signature - if !verifySignature(c.config.Token(), msgSignature, timestamp, nonce, echostr) { - logger.WarnCF("wecom_app", "Signature verification failed", map[string]any{ - "token": c.config.Token(), - "msg_signature": msgSignature, - "timestamp": timestamp, - "nonce": nonce, - }) - http.Error(w, "Invalid signature", http.StatusForbidden) - return - } - - logger.DebugC("wecom_app", "Signature verification passed") - - // Decrypt echostr with CorpID verification - // For WeCom App (自建应用), receiveid should be corp_id - logger.DebugCF("wecom_app", "Attempting to decrypt echostr", map[string]any{ - "encoding_aes_key": c.config.EncodingAESKey(), - "corp_id": c.config.CorpID, - }) - decryptedEchoStr, err := decryptMessageWithVerify(echostr, c.config.EncodingAESKey(), c.config.CorpID) - if err != nil { - logger.ErrorCF("wecom_app", "Failed to decrypt echostr", map[string]any{ - "error": err.Error(), - "encoding_aes_key": c.config.EncodingAESKey, - "corp_id": c.config.CorpID, - }) - http.Error(w, "Decryption failed", http.StatusInternalServerError) - return - } - - logger.DebugCF("wecom_app", "Successfully decrypted echostr", map[string]any{ - "decrypted": decryptedEchoStr, - }) - - // Remove BOM and whitespace as per WeCom documentation - // The response must be plain text without quotes, BOM, or newlines - decryptedEchoStr = strings.TrimSpace(decryptedEchoStr) - decryptedEchoStr = strings.TrimPrefix(decryptedEchoStr, "\xef\xbb\xbf") // Remove UTF-8 BOM - w.Write([]byte(decryptedEchoStr)) -} - -// handleMessageCallback handles incoming messages from WeCom -func (c *WeComAppChannel) handleMessageCallback(ctx context.Context, w http.ResponseWriter, r *http.Request) { - query := r.URL.Query() - msgSignature := query.Get("msg_signature") - timestamp := query.Get("timestamp") - nonce := query.Get("nonce") - - if msgSignature == "" || timestamp == "" || nonce == "" { - http.Error(w, "Missing parameters", http.StatusBadRequest) - return - } - - // Read request body - body, err := io.ReadAll(r.Body) - if err != nil { - http.Error(w, "Failed to read body", http.StatusBadRequest) - return - } - defer r.Body.Close() - - // Parse XML to get encrypted message - var encryptedMsg struct { - XMLName xml.Name `xml:"xml"` - ToUserName string `xml:"ToUserName"` - Encrypt string `xml:"Encrypt"` - AgentID string `xml:"AgentID"` - } - - if err = xml.Unmarshal(body, &encryptedMsg); err != nil { - logger.ErrorCF("wecom_app", "Failed to parse XML", map[string]any{ - "error": err.Error(), - }) - http.Error(w, "Invalid XML", http.StatusBadRequest) - return - } - - // Verify signature - if !verifySignature(c.config.Token(), msgSignature, timestamp, nonce, encryptedMsg.Encrypt) { - logger.WarnC("wecom_app", "Message signature verification failed") - http.Error(w, "Invalid signature", http.StatusForbidden) - return - } - - // Decrypt message with CorpID verification - // For WeCom App (自建应用), receiveid should be corp_id - decryptedMsg, err := decryptMessageWithVerify(encryptedMsg.Encrypt, c.config.EncodingAESKey(), c.config.CorpID) - if err != nil { - logger.ErrorCF("wecom_app", "Failed to decrypt message", map[string]any{ - "error": err.Error(), - }) - http.Error(w, "Decryption failed", http.StatusInternalServerError) - return - } - - // Parse decrypted XML message - var msg WeComXMLMessage - if err := xml.Unmarshal([]byte(decryptedMsg), &msg); err != nil { - logger.ErrorCF("wecom_app", "Failed to parse decrypted message", map[string]any{ - "error": err.Error(), - }) - http.Error(w, "Invalid message format", http.StatusBadRequest) - return - } - - // Process the message with the channel's long-lived context (not the HTTP - // request context, which is canceled as soon as we return the response). - go c.processMessage(c.ctx, msg) - - // Return success response immediately - // WeCom App requires response within configured timeout (default 5 seconds) - w.Write([]byte("success")) -} - -// processMessage processes the received message -func (c *WeComAppChannel) processMessage(ctx context.Context, msg WeComXMLMessage) { - // Skip non-text messages for now (can be extended) - if msg.MsgType != "text" && msg.MsgType != "image" && msg.MsgType != "voice" { - logger.DebugCF("wecom_app", "Skipping non-supported message type", map[string]any{ - "msg_type": msg.MsgType, - }) - return - } - - // Message deduplication: Use msg_id to prevent duplicate processing - // As per WeCom documentation, use msg_id for deduplication - msgID := fmt.Sprintf("%d", msg.MsgId) - if !c.processedMsgs.MarkMessageProcessed(msgID) { - logger.DebugCF("wecom_app", "Skipping duplicate message", map[string]any{ - "msg_id": msgID, - }) - return - } - - senderID := msg.FromUserName - chatID := senderID // WeCom App uses user ID as chat ID for direct messages - - // Build metadata - // WeCom App only supports direct messages (private chat) - peer := bus.Peer{Kind: "direct", ID: senderID} - messageID := fmt.Sprintf("%d", msg.MsgId) - - metadata := map[string]string{ - "msg_type": msg.MsgType, - "msg_id": fmt.Sprintf("%d", msg.MsgId), - "agent_id": fmt.Sprintf("%d", msg.AgentID), - "platform": "wecom_app", - "media_id": msg.MediaId, - "create_time": fmt.Sprintf("%d", msg.CreateTime), - } - - content := msg.Content - - logger.DebugCF("wecom_app", "Received message", map[string]any{ - "sender_id": senderID, - "msg_type": msg.MsgType, - "preview": utils.Truncate(content, 50), - }) - - // Build sender info - appSender := bus.SenderInfo{ - Platform: "wecom", - PlatformID: senderID, - CanonicalID: identity.BuildCanonicalID("wecom", senderID), - } - - // Handle the message through the base channel - c.HandleMessage(ctx, peer, messageID, senderID, chatID, content, nil, metadata, appSender) -} - -// tokenRefreshLoop periodically refreshes the access token -func (c *WeComAppChannel) tokenRefreshLoop() { - ticker := time.NewTicker(5 * time.Minute) - defer ticker.Stop() - - for { - select { - case <-c.ctx.Done(): - return - case <-ticker.C: - if err := c.refreshAccessToken(); err != nil { - logger.ErrorCF("wecom_app", "Failed to refresh access token", map[string]any{ - "error": err.Error(), - }) - } - } - } -} - -// refreshAccessToken gets a new access token from WeCom API -func (c *WeComAppChannel) refreshAccessToken() error { - apiURL := fmt.Sprintf("%s/cgi-bin/gettoken?corpid=%s&corpsecret=%s", - wecomAPIBase, url.QueryEscape(c.config.CorpID), url.QueryEscape(c.config.CorpSecret())) - - resp, err := http.Get(apiURL) - if err != nil { - return fmt.Errorf("failed to request access token: %w", err) - } - defer resp.Body.Close() - - body, err := io.ReadAll(resp.Body) - if err != nil { - return fmt.Errorf("failed to read response: %w", err) - } - - var tokenResp WeComAccessTokenResponse - if err := json.Unmarshal(body, &tokenResp); err != nil { - return fmt.Errorf("failed to parse response: %w", err) - } - - if tokenResp.ErrCode != 0 { - return fmt.Errorf("API error: %s (code: %d)", tokenResp.ErrMsg, tokenResp.ErrCode) - } - - c.tokenMu.Lock() - c.accessToken = tokenResp.AccessToken - c.tokenExpiry = time.Now().Add(time.Duration(tokenResp.ExpiresIn-300) * time.Second) // Refresh 5 minutes early - c.tokenMu.Unlock() - - logger.DebugC("wecom_app", "Access token refreshed successfully") - return nil -} - -// getAccessToken returns the current valid access token -func (c *WeComAppChannel) getAccessToken() string { - c.tokenMu.RLock() - defer c.tokenMu.RUnlock() - - if time.Now().After(c.tokenExpiry) { - return "" - } - - return c.accessToken -} - -// sendTextMessage sends a text message to a user. -func (c *WeComAppChannel) sendTextMessage(ctx context.Context, accessToken, userID, content string) error { - msg := WeComTextMessage{ - ToUser: userID, - MsgType: "text", - AgentID: c.config.AgentID, - } - msg.Text.Content = content - return c.sendWeComMessage(ctx, accessToken, msg) -} - -// handleHealth handles health check requests -func (c *WeComAppChannel) handleHealth(w http.ResponseWriter, r *http.Request) { - status := map[string]any{ - "status": "ok", - "running": c.IsRunning(), - "has_token": c.getAccessToken() != "", - } - - w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(status) -} diff --git a/pkg/channels/wecom/app_test.go b/pkg/channels/wecom/app_test.go deleted file mode 100644 index 502544441..000000000 --- a/pkg/channels/wecom/app_test.go +++ /dev/null @@ -1,1060 +0,0 @@ -package wecom - -import ( - "bytes" - "context" - "crypto/aes" - "crypto/cipher" - "crypto/sha1" - "encoding/base64" - "encoding/binary" - "encoding/json" - "encoding/xml" - "fmt" - "net/http" - "net/http/httptest" - "sort" - "strings" - "testing" - "time" - - "github.com/sipeed/picoclaw/pkg/bus" - "github.com/sipeed/picoclaw/pkg/config" -) - -// generateTestAESKeyApp generates a valid test AES key for WeCom App -func generateTestAESKeyApp() string { - // AES key needs to be 32 bytes (256 bits) for AES-256 - key := make([]byte, 32) - for i := range key { - key[i] = byte(i + 1) - } - // Return base64 encoded key without padding - return base64.StdEncoding.EncodeToString(key)[:43] -} - -// encryptTestMessageApp encrypts a message for testing WeCom App -func encryptTestMessageApp(message, aesKey string) (string, error) { - // Decode AES key - key, err := base64.StdEncoding.DecodeString(aesKey + "=") - if err != nil { - return "", err - } - - // Prepare message: random(16) + msg_len(4) + msg + corp_id - random := make([]byte, 0, 16) - for i := range 16 { - random = append(random, byte(i+1)) - } - - msgBytes := []byte(message) - corpID := []byte("test_corp_id") - - msgLen := uint32(len(msgBytes)) - lenBytes := make([]byte, 4) - binary.BigEndian.PutUint32(lenBytes, msgLen) - - plainText := append(random, lenBytes...) - plainText = append(plainText, msgBytes...) - plainText = append(plainText, corpID...) - - // PKCS7 padding - blockSize := aes.BlockSize - padding := blockSize - len(plainText)%blockSize - padText := bytes.Repeat([]byte{byte(padding)}, padding) - plainText = append(plainText, padText...) - - // Encrypt - block, err := aes.NewCipher(key) - if err != nil { - return "", err - } - - mode := cipher.NewCBCEncrypter(block, key[:aes.BlockSize]) - cipherText := make([]byte, len(plainText)) - mode.CryptBlocks(cipherText, plainText) - - return base64.StdEncoding.EncodeToString(cipherText), nil -} - -// generateSignatureApp generates a signature for testing WeCom App -func generateSignatureApp(token, timestamp, nonce, msgEncrypt string) string { - params := []string{token, timestamp, nonce, msgEncrypt} - sort.Strings(params) - str := strings.Join(params, "") - hash := sha1.Sum([]byte(str)) - return fmt.Sprintf("%x", hash) -} - -func TestNewWeComAppChannel(t *testing.T) { - msgBus := bus.NewMessageBus() - - t.Run("missing corp_id", func(t *testing.T) { - cfg := config.WeComAppConfig{ - CorpID: "", - AgentID: 1000002, - } - cfg.SetCorpSecret("test_secret") - _, err := NewWeComAppChannel(cfg, msgBus) - if err == nil { - t.Error("expected error for missing corp_id, got nil") - } - }) - - t.Run("missing corp_secret", func(t *testing.T) { - cfg := config.WeComAppConfig{ - CorpID: "test_corp_id", - AgentID: 1000002, - } - _, err := NewWeComAppChannel(cfg, msgBus) - if err == nil { - t.Error("expected error for missing corp_secret, got nil") - } - }) - - t.Run("missing agent_id", func(t *testing.T) { - cfg := config.WeComAppConfig{ - CorpID: "test_corp_id", - AgentID: 0, - } - cfg.SetCorpSecret("test_secret") - _, err := NewWeComAppChannel(cfg, msgBus) - if err == nil { - t.Error("expected error for missing agent_id, got nil") - } - }) - - t.Run("valid config", func(t *testing.T) { - cfg := config.WeComAppConfig{ - CorpID: "test_corp_id", - AgentID: 1000002, - AllowFrom: []string{"user1", "user2"}, - } - cfg.SetCorpSecret("test_secret") - ch, err := NewWeComAppChannel(cfg, msgBus) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - if ch.Name() != "wecom_app" { - t.Errorf("Name() = %q, want %q", ch.Name(), "wecom_app") - } - if ch.IsRunning() { - t.Error("new channel should not be running") - } - }) -} - -func TestWeComAppChannelIsAllowed(t *testing.T) { - msgBus := bus.NewMessageBus() - - t.Run("empty allowlist allows all", func(t *testing.T) { - cfg := config.WeComAppConfig{ - CorpID: "test_corp_id", - AgentID: 1000002, - AllowFrom: []string{}, - } - cfg.SetCorpSecret("test_secret") - ch, _ := NewWeComAppChannel(cfg, msgBus) - if !ch.IsAllowed("any_user") { - t.Error("empty allowlist should allow all users") - } - }) - - t.Run("allowlist restricts users", func(t *testing.T) { - cfg := config.WeComAppConfig{ - CorpID: "test_corp_id", - AgentID: 1000002, - AllowFrom: []string{"allowed_user"}, - } - cfg.SetCorpSecret("test_secret") - ch, _ := NewWeComAppChannel(cfg, msgBus) - if !ch.IsAllowed("allowed_user") { - t.Error("allowed user should pass allowlist check") - } - if ch.IsAllowed("blocked_user") { - t.Error("non-allowed user should be blocked") - } - }) -} - -func TestWeComAppVerifySignature(t *testing.T) { - msgBus := bus.NewMessageBus() - cfg := config.WeComAppConfig{} - cfg.CorpID = "test_corp_id" - cfg.SetCorpSecret("test_secret") - cfg.AgentID = 1000002 - cfg.SetToken("test_token") - ch, _ := NewWeComAppChannel(cfg, msgBus) - - t.Run("valid signature", func(t *testing.T) { - timestamp := "1234567890" - nonce := "test_nonce" - msgEncrypt := "test_message" - expectedSig := generateSignatureApp("test_token", timestamp, nonce, msgEncrypt) - - if !verifySignature(ch.config.Token(), expectedSig, timestamp, nonce, msgEncrypt) { - t.Error("valid signature should pass verification") - } - }) - - t.Run("invalid signature", func(t *testing.T) { - timestamp := "1234567890" - nonce := "test_nonce" - msgEncrypt := "test_message" - - if verifySignature(ch.config.Token(), "invalid_sig", timestamp, nonce, msgEncrypt) { - t.Error("invalid signature should fail verification") - } - }) - - t.Run("empty token rejects verification (fail-closed)", func(t *testing.T) { - cfgEmpty := config.WeComAppConfig{} - cfgEmpty.CorpID = "test_corp_id" - cfgEmpty.SetCorpSecret("test_secret") - cfgEmpty.AgentID = 1000002 - cfgEmpty.SetToken("") - chEmpty, _ := NewWeComAppChannel(cfgEmpty, msgBus) - - if verifySignature(chEmpty.config.Token(), "any_sig", "any_ts", "any_nonce", "any_msg") { - t.Error("empty token should reject verification (fail-closed)") - } - }) -} - -func TestWeComAppDecryptMessage(t *testing.T) { - msgBus := bus.NewMessageBus() - - t.Run("decrypt without AES key", func(t *testing.T) { - cfg := config.WeComAppConfig{} - cfg.CorpID = "test_corp_id" - cfg.SetCorpSecret("test_secret") - cfg.AgentID = 1000002 - cfg.SetEncodingAESKey("") - ch, _ := NewWeComAppChannel(cfg, msgBus) - - // Without AES key, message should be base64 decoded only - plainText := "hello world" - encoded := base64.StdEncoding.EncodeToString([]byte(plainText)) - - result, err := decryptMessage(encoded, ch.config.EncodingAESKey()) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - if result != plainText { - t.Errorf("decryptMessage() = %q, want %q", result, plainText) - } - }) - - t.Run("decrypt with AES key", func(t *testing.T) { - aesKey := generateTestAESKeyApp() - cfg := config.WeComAppConfig{ - CorpID: "test_corp_id", - AgentID: 1000002, - } - cfg.SetCorpSecret("test_secret") - cfg.SetEncodingAESKey(aesKey) - ch, _ := NewWeComAppChannel(cfg, msgBus) - - originalMsg := "Hello" - encrypted, err := encryptTestMessageApp(originalMsg, aesKey) - if err != nil { - t.Fatalf("failed to encrypt test message: %v", err) - } - - result, err := decryptMessage(encrypted, ch.config.EncodingAESKey()) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - if result != originalMsg { - t.Errorf("WeComDecryptMessage() = %q, want %q", result, originalMsg) - } - }) - - t.Run("invalid base64", func(t *testing.T) { - cfg := config.WeComAppConfig{ - CorpID: "test_corp_id", - AgentID: 1000002, - } - cfg.SetCorpSecret("test_secret") - cfg.SetEncodingAESKey("") - ch, _ := NewWeComAppChannel(cfg, msgBus) - - _, err := decryptMessage("invalid_base64!!!", ch.config.EncodingAESKey()) - if err == nil { - t.Error("expected error for invalid base64, got nil") - } - }) - - t.Run("invalid AES key", func(t *testing.T) { - cfg := config.WeComAppConfig{} - cfg.CorpID = "test_corp_id" - cfg.SetCorpSecret("test_secret") - cfg.AgentID = 1000002 - cfg.SetEncodingAESKey("invalid_key") - ch, _ := NewWeComAppChannel(cfg, msgBus) - - _, err := decryptMessage(base64.StdEncoding.EncodeToString([]byte("test")), ch.config.EncodingAESKey()) - if err == nil { - t.Error("expected error for invalid AES key, got nil") - } - }) - - t.Run("ciphertext too short", func(t *testing.T) { - aesKey := generateTestAESKeyApp() - cfg := config.WeComAppConfig{} - cfg.CorpID = "test_corp_id" - cfg.SetCorpSecret("test_secret") - cfg.AgentID = 1000002 - cfg.SetEncodingAESKey(aesKey) - ch, _ := NewWeComAppChannel(cfg, msgBus) - - // Encrypt a very short message that results in ciphertext less than block size - shortData := make([]byte, 8) - _, err := decryptMessage(base64.StdEncoding.EncodeToString(shortData), ch.config.EncodingAESKey()) - if err == nil { - t.Error("expected error for short ciphertext, got nil") - } - }) -} - -func TestWeComAppHandleVerification(t *testing.T) { - msgBus := bus.NewMessageBus() - aesKey := generateTestAESKeyApp() - cfg := config.WeComAppConfig{} - cfg.CorpID = "test_corp_id" - cfg.SetCorpSecret("test_secret") - cfg.AgentID = 1000002 - cfg.SetToken("test_token") - cfg.SetEncodingAESKey(aesKey) - ch, _ := NewWeComAppChannel(cfg, msgBus) - - t.Run("valid verification request", func(t *testing.T) { - echostr := "test_echostr_123" - encryptedEchostr, _ := encryptTestMessageApp(echostr, aesKey) - timestamp := "1234567890" - nonce := "test_nonce" - signature := generateSignatureApp("test_token", timestamp, nonce, encryptedEchostr) - - req := httptest.NewRequest( - http.MethodGet, - "/webhook/wecom-app?msg_signature="+signature+"×tamp="+timestamp+"&nonce="+nonce+"&echostr="+encryptedEchostr, - nil, - ) - w := httptest.NewRecorder() - - ch.handleVerification(context.Background(), w, req) - - if w.Code != http.StatusOK { - t.Errorf("status code = %d, want %d", w.Code, http.StatusOK) - } - if w.Body.String() != echostr { - t.Errorf("response body = %q, want %q", w.Body.String(), echostr) - } - }) - - t.Run("missing parameters", func(t *testing.T) { - req := httptest.NewRequest(http.MethodGet, "/webhook/wecom-app?msg_signature=sig×tamp=ts", nil) - w := httptest.NewRecorder() - - ch.handleVerification(context.Background(), w, req) - - if w.Code != http.StatusBadRequest { - t.Errorf("status code = %d, want %d", w.Code, http.StatusBadRequest) - } - }) - - t.Run("invalid signature", func(t *testing.T) { - echostr := "test_echostr" - encryptedEchostr, _ := encryptTestMessageApp(echostr, aesKey) - timestamp := "1234567890" - nonce := "test_nonce" - - req := httptest.NewRequest( - http.MethodGet, - "/webhook/wecom-app?msg_signature=invalid_sig×tamp="+timestamp+"&nonce="+nonce+"&echostr="+encryptedEchostr, - nil, - ) - w := httptest.NewRecorder() - - ch.handleVerification(context.Background(), w, req) - - if w.Code != http.StatusForbidden { - t.Errorf("status code = %d, want %d", w.Code, http.StatusForbidden) - } - }) -} - -func TestWeComAppHandleMessageCallback(t *testing.T) { - msgBus := bus.NewMessageBus() - aesKey := generateTestAESKeyApp() - cfg := config.WeComAppConfig{} - cfg.CorpID = "test_corp_id" - cfg.SetCorpSecret("test_secret") - cfg.AgentID = 1000002 - cfg.SetToken("test_token") - cfg.SetEncodingAESKey(aesKey) - ch, _ := NewWeComAppChannel(cfg, msgBus) - - t.Run("valid message callback", func(t *testing.T) { - // Create XML message - xmlMsg := WeComXMLMessage{ - ToUserName: "corp_id", - FromUserName: "user123", - CreateTime: 1234567890, - MsgType: "text", - Content: "Hello World", - MsgId: 123456, - AgentID: 1000002, - } - xmlData, _ := xml.Marshal(xmlMsg) - - // Encrypt message - encrypted, _ := encryptTestMessageApp(string(xmlData), aesKey) - - // Create encrypted XML wrapper - encryptedWrapper := struct { - XMLName xml.Name `xml:"xml"` - Encrypt string `xml:"Encrypt"` - }{ - Encrypt: encrypted, - } - wrapperData, _ := xml.Marshal(encryptedWrapper) - - timestamp := "1234567890" - nonce := "test_nonce" - signature := generateSignatureApp("test_token", timestamp, nonce, encrypted) - - req := httptest.NewRequest( - http.MethodPost, - "/webhook/wecom-app?msg_signature="+signature+"×tamp="+timestamp+"&nonce="+nonce, - bytes.NewReader(wrapperData), - ) - w := httptest.NewRecorder() - - ch.handleMessageCallback(context.Background(), w, req) - - if w.Code != http.StatusOK { - t.Errorf("status code = %d, want %d", w.Code, http.StatusOK) - } - if w.Body.String() != "success" { - t.Errorf("response body = %q, want %q", w.Body.String(), "success") - } - }) - - t.Run("missing parameters", func(t *testing.T) { - req := httptest.NewRequest(http.MethodPost, "/webhook/wecom-app?msg_signature=sig", nil) - w := httptest.NewRecorder() - - ch.handleMessageCallback(context.Background(), w, req) - - if w.Code != http.StatusBadRequest { - t.Errorf("status code = %d, want %d", w.Code, http.StatusBadRequest) - } - }) - - t.Run("invalid XML", func(t *testing.T) { - timestamp := "1234567890" - nonce := "test_nonce" - signature := generateSignatureApp("test_token", timestamp, nonce, "") - - req := httptest.NewRequest( - http.MethodPost, - "/webhook/wecom-app?msg_signature="+signature+"×tamp="+timestamp+"&nonce="+nonce, - strings.NewReader("invalid xml"), - ) - w := httptest.NewRecorder() - - ch.handleMessageCallback(context.Background(), w, req) - - if w.Code != http.StatusBadRequest { - t.Errorf("status code = %d, want %d", w.Code, http.StatusBadRequest) - } - }) - - t.Run("invalid signature", func(t *testing.T) { - encryptedWrapper := struct { - XMLName xml.Name `xml:"xml"` - Encrypt string `xml:"Encrypt"` - }{ - Encrypt: "encrypted_data", - } - wrapperData, _ := xml.Marshal(encryptedWrapper) - - timestamp := "1234567890" - nonce := "test_nonce" - - req := httptest.NewRequest( - http.MethodPost, - "/webhook/wecom-app?msg_signature=invalid_sig×tamp="+timestamp+"&nonce="+nonce, - bytes.NewReader(wrapperData), - ) - w := httptest.NewRecorder() - - ch.handleMessageCallback(context.Background(), w, req) - - if w.Code != http.StatusForbidden { - t.Errorf("status code = %d, want %d", w.Code, http.StatusForbidden) - } - }) -} - -func TestWeComAppProcessMessage(t *testing.T) { - msgBus := bus.NewMessageBus() - cfg := config.WeComAppConfig{ - CorpID: "test_corp_id", - AgentID: 1000002, - } - cfg.SetCorpSecret("test_secret") - ch, _ := NewWeComAppChannel(cfg, msgBus) - - t.Run("process text message", func(t *testing.T) { - msg := WeComXMLMessage{ - ToUserName: "corp_id", - FromUserName: "user123", - CreateTime: 1234567890, - MsgType: "text", - Content: "Hello World", - MsgId: 123456, - AgentID: 1000002, - } - - // Should not panic - ch.processMessage(context.Background(), msg) - }) - - t.Run("process image message", func(t *testing.T) { - msg := WeComXMLMessage{ - ToUserName: "corp_id", - FromUserName: "user123", - CreateTime: 1234567890, - MsgType: "image", - PicUrl: "https://example.com/image.jpg", - MediaId: "media_123", - MsgId: 123456, - AgentID: 1000002, - } - - // Should not panic - ch.processMessage(context.Background(), msg) - }) - - t.Run("process voice message", func(t *testing.T) { - msg := WeComXMLMessage{ - ToUserName: "corp_id", - FromUserName: "user123", - CreateTime: 1234567890, - MsgType: "voice", - MediaId: "media_123", - Format: "amr", - MsgId: 123456, - AgentID: 1000002, - } - - // Should not panic - ch.processMessage(context.Background(), msg) - }) - - t.Run("skip unsupported message type", func(t *testing.T) { - msg := WeComXMLMessage{ - ToUserName: "corp_id", - FromUserName: "user123", - CreateTime: 1234567890, - MsgType: "video", - MsgId: 123456, - AgentID: 1000002, - } - - // Should not panic - ch.processMessage(context.Background(), msg) - }) - - t.Run("process event message", func(t *testing.T) { - msg := WeComXMLMessage{ - ToUserName: "corp_id", - FromUserName: "user123", - CreateTime: 1234567890, - MsgType: "event", - Event: "subscribe", - MsgId: 123456, - AgentID: 1000002, - } - - // Should not panic - ch.processMessage(context.Background(), msg) - }) -} - -func TestWeComAppHandleWebhook(t *testing.T) { - msgBus := bus.NewMessageBus() - cfg := config.WeComAppConfig{} - cfg.CorpID = "test_corp_id" - cfg.SetCorpSecret("test_secret") - cfg.AgentID = 1000002 - cfg.SetToken("test_token") - ch, _ := NewWeComAppChannel(cfg, msgBus) - - t.Run("GET request calls verification", func(t *testing.T) { - echostr := "test_echostr" - encoded := base64.StdEncoding.EncodeToString([]byte(echostr)) - timestamp := "1234567890" - nonce := "test_nonce" - signature := generateSignatureApp("test_token", timestamp, nonce, encoded) - - req := httptest.NewRequest( - http.MethodGet, - "/webhook/wecom-app?msg_signature="+signature+"×tamp="+timestamp+"&nonce="+nonce+"&echostr="+encoded, - nil, - ) - w := httptest.NewRecorder() - - ch.handleWebhook(w, req) - - if w.Code != http.StatusOK { - t.Errorf("status code = %d, want %d", w.Code, http.StatusOK) - } - }) - - t.Run("POST request calls message callback", func(t *testing.T) { - encryptedWrapper := struct { - XMLName xml.Name `xml:"xml"` - Encrypt string `xml:"Encrypt"` - }{ - Encrypt: base64.StdEncoding.EncodeToString([]byte("test")), - } - wrapperData, _ := xml.Marshal(encryptedWrapper) - - timestamp := "1234567890" - nonce := "test_nonce" - signature := generateSignatureApp("test_token", timestamp, nonce, encryptedWrapper.Encrypt) - - req := httptest.NewRequest( - http.MethodPost, - "/webhook/wecom-app?msg_signature="+signature+"×tamp="+timestamp+"&nonce="+nonce, - bytes.NewReader(wrapperData), - ) - w := httptest.NewRecorder() - - ch.handleWebhook(w, req) - - // Should not be method not allowed - if w.Code == http.StatusMethodNotAllowed { - t.Error("POST request should not return Method Not Allowed") - } - }) - - t.Run("unsupported method", func(t *testing.T) { - req := httptest.NewRequest(http.MethodPut, "/webhook/wecom-app", nil) - w := httptest.NewRecorder() - - ch.handleWebhook(w, req) - - if w.Code != http.StatusMethodNotAllowed { - t.Errorf("status code = %d, want %d", w.Code, http.StatusMethodNotAllowed) - } - }) -} - -func TestWeComAppHandleHealth(t *testing.T) { - msgBus := bus.NewMessageBus() - cfg := config.WeComAppConfig{ - CorpID: "test_corp_id", - AgentID: 1000002, - } - cfg.SetCorpSecret("test_secret") - ch, _ := NewWeComAppChannel(cfg, msgBus) - - req := httptest.NewRequest(http.MethodGet, "/health/wecom-app", nil) - w := httptest.NewRecorder() - - ch.handleHealth(w, req) - - if w.Code != http.StatusOK { - t.Errorf("status code = %d, want %d", w.Code, http.StatusOK) - } - - contentType := w.Header().Get("Content-Type") - if contentType != "application/json" { - t.Errorf("Content-Type = %q, want %q", contentType, "application/json") - } - - body := w.Body.String() - if !strings.Contains(body, "status") || !strings.Contains(body, "running") || !strings.Contains(body, "has_token") { - t.Errorf("response body should contain status, running, and has_token fields, got: %s", body) - } -} - -func TestWeComAppAccessToken(t *testing.T) { - msgBus := bus.NewMessageBus() - cfg := config.WeComAppConfig{ - CorpID: "test_corp_id", - AgentID: 1000002, - } - cfg.SetCorpSecret("test_secret") - ch, _ := NewWeComAppChannel(cfg, msgBus) - - t.Run("get empty access token initially", func(t *testing.T) { - token := ch.getAccessToken() - if token != "" { - t.Errorf("getAccessToken() = %q, want empty string", token) - } - }) - - t.Run("set and get access token", func(t *testing.T) { - ch.tokenMu.Lock() - ch.accessToken = "test_token_123" - ch.tokenExpiry = time.Now().Add(1 * time.Hour) - ch.tokenMu.Unlock() - - token := ch.getAccessToken() - if token != "test_token_123" { - t.Errorf("getAccessToken() = %q, want %q", token, "test_token_123") - } - }) - - t.Run("expired token returns empty", func(t *testing.T) { - ch.tokenMu.Lock() - ch.accessToken = "expired_token" - ch.tokenExpiry = time.Now().Add(-1 * time.Hour) - ch.tokenMu.Unlock() - - token := ch.getAccessToken() - if token != "" { - t.Errorf("getAccessToken() = %q, want empty string for expired token", token) - } - }) -} - -func TestWeComAppMessageStructures(t *testing.T) { - t.Run("WeComTextMessage structure", func(t *testing.T) { - msg := WeComTextMessage{ - ToUser: "user123", - MsgType: "text", - AgentID: 1000002, - } - msg.Text.Content = "Hello World" - - if msg.ToUser != "user123" { - t.Errorf("ToUser = %q, want %q", msg.ToUser, "user123") - } - if msg.MsgType != "text" { - t.Errorf("MsgType = %q, want %q", msg.MsgType, "text") - } - if msg.AgentID != 1000002 { - t.Errorf("AgentID = %d, want %d", msg.AgentID, 1000002) - } - if msg.Text.Content != "Hello World" { - t.Errorf("Text.Content = %q, want %q", msg.Text.Content, "Hello World") - } - - // Test JSON marshaling - jsonData, err := json.Marshal(msg) - if err != nil { - t.Fatalf("failed to marshal JSON: %v", err) - } - - var unmarshaled WeComTextMessage - err = json.Unmarshal(jsonData, &unmarshaled) - if err != nil { - t.Fatalf("failed to unmarshal JSON: %v", err) - } - - if unmarshaled.ToUser != msg.ToUser { - t.Errorf("JSON round-trip failed for ToUser") - } - }) - - t.Run("WeComMarkdownMessage structure", func(t *testing.T) { - msg := WeComMarkdownMessage{ - ToUser: "user123", - MsgType: "markdown", - AgentID: 1000002, - } - msg.Markdown.Content = "# Hello\nWorld" - - if msg.Markdown.Content != "# Hello\nWorld" { - t.Errorf("Markdown.Content = %q, want %q", msg.Markdown.Content, "# Hello\nWorld") - } - - // Test JSON marshaling - jsonData, err := json.Marshal(msg) - if err != nil { - t.Fatalf("failed to marshal JSON: %v", err) - } - - if !bytes.Contains(jsonData, []byte("markdown")) { - t.Error("JSON should contain 'markdown' field") - } - }) - - t.Run("WeComImageMessage structure", func(t *testing.T) { - msg := WeComImageMessage{ - ToUser: "user123", - MsgType: "image", - AgentID: 1000002, - } - msg.Image.MediaID = "media_123456" - - if msg.Image.MediaID != "media_123456" { - t.Errorf("Image.MediaID = %q, want %q", msg.Image.MediaID, "media_123456") - } - if msg.ToUser != "user123" { - t.Errorf("ToUser = %q, want %q", msg.ToUser, "user123") - } - if msg.MsgType != "image" { - t.Errorf("MsgType = %q, want %q", msg.MsgType, "image") - } - if msg.AgentID != 1000002 { - t.Errorf("AgentID = %d, want %d", msg.AgentID, 1000002) - } - }) - - t.Run("WeComAccessTokenResponse structure", func(t *testing.T) { - jsonData := `{ - "errcode": 0, - "errmsg": "ok", - "access_token": "test_access_token", - "expires_in": 7200 - }` - - var resp WeComAccessTokenResponse - err := json.Unmarshal([]byte(jsonData), &resp) - if err != nil { - t.Fatalf("failed to unmarshal JSON: %v", err) - } - - if resp.ErrCode != 0 { - t.Errorf("ErrCode = %d, want %d", resp.ErrCode, 0) - } - if resp.ErrMsg != "ok" { - t.Errorf("ErrMsg = %q, want %q", resp.ErrMsg, "ok") - } - if resp.AccessToken != "test_access_token" { - t.Errorf("AccessToken = %q, want %q", resp.AccessToken, "test_access_token") - } - if resp.ExpiresIn != 7200 { - t.Errorf("ExpiresIn = %d, want %d", resp.ExpiresIn, 7200) - } - }) - - t.Run("WeComSendMessageResponse structure", func(t *testing.T) { - jsonData := `{ - "errcode": 0, - "errmsg": "ok", - "invaliduser": "", - "invalidparty": "", - "invalidtag": "" - }` - - var resp WeComSendMessageResponse - err := json.Unmarshal([]byte(jsonData), &resp) - if err != nil { - t.Fatalf("failed to unmarshal JSON: %v", err) - } - - if resp.ErrCode != 0 { - t.Errorf("ErrCode = %d, want %d", resp.ErrCode, 0) - } - if resp.ErrMsg != "ok" { - t.Errorf("ErrMsg = %q, want %q", resp.ErrMsg, "ok") - } - }) -} - -func TestWeComAppXMLMessageStructure(t *testing.T) { - xmlData := ` - - - - 1234567890 - - - 1234567890123456 - 1000002 -` - - var msg WeComXMLMessage - err := xml.Unmarshal([]byte(xmlData), &msg) - if err != nil { - t.Fatalf("failed to unmarshal XML: %v", err) - } - - if msg.ToUserName != "corp_id" { - t.Errorf("ToUserName = %q, want %q", msg.ToUserName, "corp_id") - } - if msg.FromUserName != "user123" { - t.Errorf("FromUserName = %q, want %q", msg.FromUserName, "user123") - } - if msg.CreateTime != 1234567890 { - t.Errorf("CreateTime = %d, want %d", msg.CreateTime, 1234567890) - } - if msg.MsgType != "text" { - t.Errorf("MsgType = %q, want %q", msg.MsgType, "text") - } - if msg.Content != "Hello World" { - t.Errorf("Content = %q, want %q", msg.Content, "Hello World") - } - if msg.MsgId != 1234567890123456 { - t.Errorf("MsgId = %d, want %d", msg.MsgId, 1234567890123456) - } - if msg.AgentID != 1000002 { - t.Errorf("AgentID = %d, want %d", msg.AgentID, 1000002) - } -} - -func TestWeComAppXMLMessageImage(t *testing.T) { - xmlData := ` - - - - 1234567890 - - - - 1234567890123456 - 1000002 -` - - var msg WeComXMLMessage - err := xml.Unmarshal([]byte(xmlData), &msg) - if err != nil { - t.Fatalf("failed to unmarshal XML: %v", err) - } - - if msg.MsgType != "image" { - t.Errorf("MsgType = %q, want %q", msg.MsgType, "image") - } - if msg.PicUrl != "https://example.com/image.jpg" { - t.Errorf("PicUrl = %q, want %q", msg.PicUrl, "https://example.com/image.jpg") - } - if msg.MediaId != "media_123" { - t.Errorf("MediaId = %q, want %q", msg.MediaId, "media_123") - } -} - -func TestWeComAppXMLMessageVoice(t *testing.T) { - xmlData := ` - - - - 1234567890 - - - - 1234567890123456 - 1000002 -` - - var msg WeComXMLMessage - err := xml.Unmarshal([]byte(xmlData), &msg) - if err != nil { - t.Fatalf("failed to unmarshal XML: %v", err) - } - - if msg.MsgType != "voice" { - t.Errorf("MsgType = %q, want %q", msg.MsgType, "voice") - } - if msg.Format != "amr" { - t.Errorf("Format = %q, want %q", msg.Format, "amr") - } -} - -func TestWeComAppXMLMessageLocation(t *testing.T) { - xmlData := ` - - - - 1234567890 - - 39.9042 - 116.4074 - 16 - - 1234567890123456 - 1000002 -` - - var msg WeComXMLMessage - err := xml.Unmarshal([]byte(xmlData), &msg) - if err != nil { - t.Fatalf("failed to unmarshal XML: %v", err) - } - - if msg.MsgType != "location" { - t.Errorf("MsgType = %q, want %q", msg.MsgType, "location") - } - if msg.LocationX != 39.9042 { - t.Errorf("LocationX = %f, want %f", msg.LocationX, 39.9042) - } - if msg.LocationY != 116.4074 { - t.Errorf("LocationY = %f, want %f", msg.LocationY, 116.4074) - } - if msg.Scale != 16 { - t.Errorf("Scale = %d, want %d", msg.Scale, 16) - } - if msg.Label != "Beijing" { - t.Errorf("Label = %q, want %q", msg.Label, "Beijing") - } -} - -func TestWeComAppXMLMessageLink(t *testing.T) { - xmlData := ` - - - - 1234567890 - - <![CDATA[Link Title]]> - - - 1234567890123456 - 1000002 -` - - var msg WeComXMLMessage - err := xml.Unmarshal([]byte(xmlData), &msg) - if err != nil { - t.Fatalf("failed to unmarshal XML: %v", err) - } - - if msg.MsgType != "link" { - t.Errorf("MsgType = %q, want %q", msg.MsgType, "link") - } - if msg.Title != "Link Title" { - t.Errorf("Title = %q, want %q", msg.Title, "Link Title") - } - if msg.Description != "Link Description" { - t.Errorf("Description = %q, want %q", msg.Description, "Link Description") - } - if msg.Url != "https://example.com" { - t.Errorf("Url = %q, want %q", msg.Url, "https://example.com") - } -} - -func TestWeComAppXMLMessageEvent(t *testing.T) { - xmlData := ` - - - - 1234567890 - - - - 1000002 -` - - var msg WeComXMLMessage - err := xml.Unmarshal([]byte(xmlData), &msg) - if err != nil { - t.Fatalf("failed to unmarshal XML: %v", err) - } - - if msg.MsgType != "event" { - t.Errorf("MsgType = %q, want %q", msg.MsgType, "event") - } - if msg.Event != "subscribe" { - t.Errorf("Event = %q, want %q", msg.Event, "subscribe") - } - if msg.EventKey != "event_key_123" { - t.Errorf("EventKey = %q, want %q", msg.EventKey, "event_key_123") - } -} diff --git a/pkg/channels/wecom/bot.go b/pkg/channels/wecom/bot.go deleted file mode 100644 index 22461b768..000000000 --- a/pkg/channels/wecom/bot.go +++ /dev/null @@ -1,499 +0,0 @@ -package wecom - -import ( - "bytes" - "context" - "encoding/json" - "encoding/xml" - "fmt" - "io" - "net/http" - "strings" - "time" - - "github.com/sipeed/picoclaw/pkg/bus" - "github.com/sipeed/picoclaw/pkg/channels" - "github.com/sipeed/picoclaw/pkg/config" - "github.com/sipeed/picoclaw/pkg/identity" - "github.com/sipeed/picoclaw/pkg/logger" - "github.com/sipeed/picoclaw/pkg/utils" -) - -// WeComBotChannel implements the Channel interface for WeCom Bot (企业微信智能机器人) -// Uses webhook callback mode - simpler than WeCom App but only supports passive replies -type WeComBotChannel struct { - *channels.BaseChannel - config config.WeComConfig - client *http.Client - ctx context.Context - cancel context.CancelFunc - processedMsgs *MessageDeduplicator -} - -// WeComBotMessage represents the JSON message structure from WeCom Bot (AIBOT) -type WeComBotMessage struct { - MsgID string `json:"msgid"` - AIBotID string `json:"aibotid"` - ChatID string `json:"chatid"` // Session ID, only present for group chats - ChatType string `json:"chattype"` // "single" for DM, "group" for group chat - From struct { - UserID string `json:"userid"` - } `json:"from"` - ResponseURL string `json:"response_url"` - MsgType string `json:"msgtype"` // text, image, voice, file, mixed - Text struct { - Content string `json:"content"` - } `json:"text"` - Image struct { - URL string `json:"url"` - } `json:"image"` - Voice struct { - Content string `json:"content"` // Voice to text content - } `json:"voice"` - File struct { - URL string `json:"url"` - } `json:"file"` - Mixed struct { - MsgItem []struct { - MsgType string `json:"msgtype"` - Text struct { - Content string `json:"content"` - } `json:"text"` - Image struct { - URL string `json:"url"` - } `json:"image"` - } `json:"msg_item"` - } `json:"mixed"` - Quote struct { - MsgType string `json:"msgtype"` - Text struct { - Content string `json:"content"` - } `json:"text"` - } `json:"quote"` -} - -// WeComBotReplyMessage represents the reply message structure -type WeComBotReplyMessage struct { - MsgType string `json:"msgtype"` - Text struct { - Content string `json:"content"` - } `json:"text,omitempty"` -} - -// NewWeComBotChannel creates a new WeCom Bot channel instance -func NewWeComBotChannel(cfg config.WeComConfig, messageBus *bus.MessageBus) (*WeComBotChannel, error) { - if cfg.Token() == "" || cfg.WebhookURL == "" { - return nil, fmt.Errorf("wecom token and webhook_url are required") - } - - base := channels.NewBaseChannel("wecom", cfg, messageBus, cfg.AllowFrom, - channels.WithMaxMessageLength(2048), - channels.WithGroupTrigger(cfg.GroupTrigger), - channels.WithReasoningChannelID(cfg.ReasoningChannelID), - ) - - // Client timeout must be >= the configured ReplyTimeout so the - // per-request context deadline is always the effective limit. - clientTimeout := 30 * time.Second - if d := time.Duration(cfg.ReplyTimeout) * time.Second; d > clientTimeout { - clientTimeout = d - } - - ctx, cancel := context.WithCancel(context.Background()) - return &WeComBotChannel{ - BaseChannel: base, - config: cfg, - client: &http.Client{Timeout: clientTimeout}, - ctx: ctx, - cancel: cancel, - processedMsgs: NewMessageDeduplicator(wecomMaxProcessedMessages), - }, nil -} - -// Name returns the channel name -func (c *WeComBotChannel) Name() string { - return "wecom" -} - -// Start initializes the WeCom Bot channel -func (c *WeComBotChannel) Start(ctx context.Context) error { - logger.InfoC("wecom", "Starting WeCom Bot channel...") - - // Cancel the context created in the constructor to avoid a resource leak. - if c.cancel != nil { - c.cancel() - } - c.ctx, c.cancel = context.WithCancel(ctx) - - c.SetRunning(true) - logger.InfoC("wecom", "WeCom Bot channel started") - - return nil -} - -// Stop gracefully stops the WeCom Bot channel -func (c *WeComBotChannel) Stop(ctx context.Context) error { - logger.InfoC("wecom", "Stopping WeCom Bot channel...") - - if c.cancel != nil { - c.cancel() - } - - c.SetRunning(false) - logger.InfoC("wecom", "WeCom Bot channel stopped") - return nil -} - -// Send sends a message to WeCom user via webhook API -// Note: WeCom Bot can only reply within the configured timeout (default 5 seconds) of receiving a message -// For delayed responses, we use the webhook URL -func (c *WeComBotChannel) Send(ctx context.Context, msg bus.OutboundMessage) error { - if !c.IsRunning() { - return channels.ErrNotRunning - } - - logger.DebugCF("wecom", "Sending message via webhook", map[string]any{ - "chat_id": msg.ChatID, - "preview": utils.Truncate(msg.Content, 100), - }) - - return c.sendWebhookReply(ctx, msg.ChatID, msg.Content) -} - -// WebhookPath returns the path for registering on the shared HTTP server. -func (c *WeComBotChannel) WebhookPath() string { - if c.config.WebhookPath != "" { - return c.config.WebhookPath - } - return "/webhook/wecom" -} - -// ServeHTTP implements http.Handler for the shared HTTP server. -func (c *WeComBotChannel) ServeHTTP(w http.ResponseWriter, r *http.Request) { - c.handleWebhook(w, r) -} - -// HealthPath returns the health check endpoint path. -func (c *WeComBotChannel) HealthPath() string { - return "/health/wecom" -} - -// HealthHandler handles health check requests. -func (c *WeComBotChannel) HealthHandler(w http.ResponseWriter, r *http.Request) { - c.handleHealth(w, r) -} - -// handleWebhook handles incoming webhook requests from WeCom -func (c *WeComBotChannel) handleWebhook(w http.ResponseWriter, r *http.Request) { - ctx := r.Context() - - if r.Method == http.MethodGet { - // Handle verification request - c.handleVerification(ctx, w, r) - return - } - - if r.Method == http.MethodPost { - // Handle message callback - c.handleMessageCallback(ctx, w, r) - return - } - - http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) -} - -// handleVerification handles the URL verification request from WeCom -func (c *WeComBotChannel) handleVerification(ctx context.Context, w http.ResponseWriter, r *http.Request) { - query := r.URL.Query() - msgSignature := query.Get("msg_signature") - timestamp := query.Get("timestamp") - nonce := query.Get("nonce") - echostr := query.Get("echostr") - - if msgSignature == "" || timestamp == "" || nonce == "" || echostr == "" { - http.Error(w, "Missing parameters", http.StatusBadRequest) - return - } - - // Verify signature - if !verifySignature(c.config.Token(), msgSignature, timestamp, nonce, echostr) { - logger.WarnC("wecom", "Signature verification failed") - http.Error(w, "Invalid signature", http.StatusForbidden) - return - } - - // Decrypt echostr - // For AIBOT (智能机器人), receiveid should be empty string "" - // Reference: https://developer.work.weixin.qq.com/document/path/101033 - decryptedEchoStr, err := decryptMessageWithVerify(echostr, c.config.EncodingAESKey(), "") - if err != nil { - logger.ErrorCF("wecom", "Failed to decrypt echostr", map[string]any{ - "error": err.Error(), - }) - http.Error(w, "Decryption failed", http.StatusInternalServerError) - return - } - - // Remove BOM and whitespace as per WeCom documentation - // The response must be plain text without quotes, BOM, or newlines - decryptedEchoStr = strings.TrimSpace(decryptedEchoStr) - decryptedEchoStr = strings.TrimPrefix(decryptedEchoStr, "\xef\xbb\xbf") // Remove UTF-8 BOM - w.Write([]byte(decryptedEchoStr)) -} - -// handleMessageCallback handles incoming messages from WeCom -func (c *WeComBotChannel) handleMessageCallback(ctx context.Context, w http.ResponseWriter, r *http.Request) { - query := r.URL.Query() - msgSignature := query.Get("msg_signature") - timestamp := query.Get("timestamp") - nonce := query.Get("nonce") - - if msgSignature == "" || timestamp == "" || nonce == "" { - http.Error(w, "Missing parameters", http.StatusBadRequest) - return - } - - // Read request body - body, err := io.ReadAll(r.Body) - if err != nil { - http.Error(w, "Failed to read body", http.StatusBadRequest) - return - } - defer r.Body.Close() - - // Parse XML to get encrypted message - var encryptedMsg struct { - XMLName xml.Name `xml:"xml"` - ToUserName string `xml:"ToUserName"` - Encrypt string `xml:"Encrypt"` - AgentID string `xml:"AgentID"` - } - - if err = xml.Unmarshal(body, &encryptedMsg); err != nil { - logger.ErrorCF("wecom", "Failed to parse XML", map[string]any{ - "error": err.Error(), - }) - http.Error(w, "Invalid XML", http.StatusBadRequest) - return - } - - // Verify signature - if !verifySignature(c.config.Token(), msgSignature, timestamp, nonce, encryptedMsg.Encrypt) { - logger.WarnC("wecom", "Message signature verification failed") - http.Error(w, "Invalid signature", http.StatusForbidden) - return - } - - // Decrypt message - // For AIBOT (智能机器人), receiveid should be empty string "" - // Reference: https://developer.work.weixin.qq.com/document/path/101033 - decryptedMsg, err := decryptMessageWithVerify(encryptedMsg.Encrypt, c.config.EncodingAESKey(), "") - if err != nil { - logger.ErrorCF("wecom", "Failed to decrypt message", map[string]any{ - "error": err.Error(), - }) - http.Error(w, "Decryption failed", http.StatusInternalServerError) - return - } - - // Parse decrypted JSON message (AIBOT uses JSON format) - var msg WeComBotMessage - if err := json.Unmarshal([]byte(decryptedMsg), &msg); err != nil { - logger.ErrorCF("wecom", "Failed to parse decrypted message", map[string]any{ - "error": err.Error(), - }) - http.Error(w, "Invalid message format", http.StatusBadRequest) - return - } - - // Process the message with the channel's long-lived context (not the HTTP - // request context, which is canceled as soon as we return the response). - go c.processMessage(c.ctx, msg) - - // Return success response immediately - // WeCom Bot requires response within configured timeout (default 5 seconds) - w.Write([]byte("success")) -} - -// processMessage processes the received message -func (c *WeComBotChannel) processMessage(ctx context.Context, msg WeComBotMessage) { - // Skip unsupported message types - if msg.MsgType != "text" && msg.MsgType != "image" && msg.MsgType != "voice" && msg.MsgType != "file" && - msg.MsgType != "mixed" { - logger.DebugCF("wecom", "Skipping non-supported message type", map[string]any{ - "msg_type": msg.MsgType, - }) - return - } - - // Message deduplication: Use msg_id to prevent duplicate processing - msgID := msg.MsgID - if !c.processedMsgs.MarkMessageProcessed(msgID) { - logger.DebugCF("wecom", "Skipping duplicate message", map[string]any{ - "msg_id": msgID, - }) - return - } - - senderID := msg.From.UserID - - // Determine if this is a group chat or direct message - // ChatType: "single" for DM, "group" for group chat - isGroupChat := msg.ChatType == "group" - - var chatID, peerKind, peerID string - if isGroupChat { - // Group chat: use ChatID as chatID and peer_id - chatID = msg.ChatID - peerKind = "group" - peerID = msg.ChatID - } else { - // Direct message: use senderID as chatID and peer_id - chatID = senderID - peerKind = "direct" - peerID = senderID - } - - // Extract content based on message type - var content string - switch msg.MsgType { - case "text": - content = msg.Text.Content - case "voice": - content = msg.Voice.Content // Voice to text content - case "mixed": - // For mixed messages, concatenate text items - for _, item := range msg.Mixed.MsgItem { - if item.MsgType == "text" { - content += item.Text.Content - } - } - case "image", "file": - // For image and file, we don't have text content - content = "" - } - - // Build metadata - peer := bus.Peer{Kind: peerKind, ID: peerID} - - // In group chats, apply unified group trigger filtering - if isGroupChat { - respond, cleaned := c.ShouldRespondInGroup(false, content) - if !respond { - return - } - content = cleaned - } - - metadata := map[string]string{ - "msg_type": msg.MsgType, - "msg_id": msg.MsgID, - "platform": "wecom", - "response_url": msg.ResponseURL, - } - if isGroupChat { - metadata["chat_id"] = msg.ChatID - metadata["sender_id"] = senderID - } - - logger.DebugCF("wecom", "Received message", map[string]any{ - "sender_id": senderID, - "msg_type": msg.MsgType, - "peer_kind": peerKind, - "is_group_chat": isGroupChat, - "preview": utils.Truncate(content, 50), - }) - - // Build sender info - sender := bus.SenderInfo{ - Platform: "wecom", - PlatformID: senderID, - CanonicalID: identity.BuildCanonicalID("wecom", senderID), - } - - if !c.IsAllowedSender(sender) { - return - } - - // Handle the message through the base channel - c.HandleMessage(ctx, peer, msg.MsgID, senderID, chatID, content, nil, metadata, sender) -} - -// sendWebhookReply sends a reply using the webhook URL -func (c *WeComBotChannel) sendWebhookReply(ctx context.Context, userID, content string) error { - reply := WeComBotReplyMessage{ - MsgType: "text", - } - reply.Text.Content = content - - jsonData, err := json.Marshal(reply) - if err != nil { - return fmt.Errorf("failed to marshal reply: %w", err) - } - - // Use configurable timeout (default 5 seconds) - timeout := c.config.ReplyTimeout - if timeout <= 0 { - timeout = 5 - } - - reqCtx, cancel := context.WithTimeout(ctx, time.Duration(timeout)*time.Second) - defer cancel() - - req, err := http.NewRequestWithContext(reqCtx, http.MethodPost, c.config.WebhookURL, bytes.NewBuffer(jsonData)) - if err != nil { - return fmt.Errorf("failed to create request: %w", err) - } - req.Header.Set("Content-Type", "application/json") - - resp, err := c.client.Do(req) - if err != nil { - return channels.ClassifyNetError(err) - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - body, readErr := io.ReadAll(resp.Body) - if readErr != nil { - return channels.ClassifySendError( - resp.StatusCode, - fmt.Errorf("reading webhook error response: %w", readErr), - ) - } - return channels.ClassifySendError( - resp.StatusCode, - fmt.Errorf("webhook API error: %s", string(body)), - ) - } - - body, err := io.ReadAll(resp.Body) - if err != nil { - return fmt.Errorf("failed to read response: %w", err) - } - - // Check response - var result struct { - ErrCode int `json:"errcode"` - ErrMsg string `json:"errmsg"` - } - if err := json.Unmarshal(body, &result); err != nil { - return fmt.Errorf("failed to parse response: %w", err) - } - - if result.ErrCode != 0 { - return fmt.Errorf("webhook API error: %s (code: %d)", result.ErrMsg, result.ErrCode) - } - - return nil -} - -// handleHealth handles health check requests -func (c *WeComBotChannel) handleHealth(w http.ResponseWriter, r *http.Request) { - status := map[string]any{ - "status": "ok", - "running": c.IsRunning(), - } - - w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(status) -} diff --git a/pkg/channels/wecom/bot_test.go b/pkg/channels/wecom/bot_test.go deleted file mode 100644 index 7b50a86f7..000000000 --- a/pkg/channels/wecom/bot_test.go +++ /dev/null @@ -1,734 +0,0 @@ -package wecom - -import ( - "bytes" - "context" - "crypto/aes" - "crypto/cipher" - "crypto/sha1" - "encoding/base64" - "encoding/binary" - "encoding/json" - "encoding/xml" - "fmt" - "net/http" - "net/http/httptest" - "sort" - "strings" - "testing" - - "github.com/sipeed/picoclaw/pkg/bus" - "github.com/sipeed/picoclaw/pkg/config" -) - -// generateTestAESKey generates a valid test AES key -func generateTestAESKey() string { - // AES key needs to be 32 bytes (256 bits) for AES-256 - key := make([]byte, 32) - for i := range key { - key[i] = byte(i) - } - // Return base64 encoded key without padding - return base64.StdEncoding.EncodeToString(key)[:43] -} - -// encryptTestMessage encrypts a message for testing (AIBOT JSON format) -func encryptTestMessage(message, aesKey string) (string, error) { - // Decode AES key - key, err := base64.StdEncoding.DecodeString(aesKey + "=") - if err != nil { - return "", err - } - - // Prepare message: random(16) + msg_len(4) + msg + receiveid - random := make([]byte, 0, 16) - for i := range 16 { - random = append(random, byte(i)) - } - - msgBytes := []byte(message) - receiveID := []byte("test_aibot_id") - - msgLen := uint32(len(msgBytes)) - lenBytes := make([]byte, 4) - binary.BigEndian.PutUint32(lenBytes, msgLen) - - plainText := append(random, lenBytes...) - plainText = append(plainText, msgBytes...) - plainText = append(plainText, receiveID...) - - // PKCS7 padding - blockSize := aes.BlockSize - padding := blockSize - len(plainText)%blockSize - padText := bytes.Repeat([]byte{byte(padding)}, padding) - plainText = append(plainText, padText...) - - // Encrypt - block, err := aes.NewCipher(key) - if err != nil { - return "", err - } - - mode := cipher.NewCBCEncrypter(block, key[:aes.BlockSize]) - cipherText := make([]byte, len(plainText)) - mode.CryptBlocks(cipherText, plainText) - - return base64.StdEncoding.EncodeToString(cipherText), nil -} - -// generateSignature generates a signature for testing -func generateSignature(token, timestamp, nonce, msgEncrypt string) string { - params := []string{token, timestamp, nonce, msgEncrypt} - sort.Strings(params) - str := strings.Join(params, "") - hash := sha1.Sum([]byte(str)) - return fmt.Sprintf("%x", hash) -} - -func TestNewWeComBotChannel(t *testing.T) { - msgBus := bus.NewMessageBus() - - t.Run("missing token", func(t *testing.T) { - cfg := config.WeComConfig{} - cfg.SetToken("") - cfg.WebhookURL = "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=test" - _, err := NewWeComBotChannel(cfg, msgBus) - if err == nil { - t.Error("expected error for missing token, got nil") - } - }) - - t.Run("missing webhook_url", func(t *testing.T) { - cfg := config.WeComConfig{} - cfg.SetToken("test_token") - cfg.WebhookURL = "" - _, err := NewWeComBotChannel(cfg, msgBus) - if err == nil { - t.Error("expected error for missing webhook_url, got nil") - } - }) - - t.Run("valid config", func(t *testing.T) { - cfg := config.WeComConfig{} - cfg.SetToken("test_token") - cfg.WebhookURL = "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=test" - cfg.AllowFrom = []string{"user1", "user2"} - ch, err := NewWeComBotChannel(cfg, msgBus) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - if ch.Name() != "wecom" { - t.Errorf("Name() = %q, want %q", ch.Name(), "wecom") - } - if ch.IsRunning() { - t.Error("new channel should not be running") - } - }) -} - -func TestWeComBotChannelIsAllowed(t *testing.T) { - msgBus := bus.NewMessageBus() - - t.Run("empty allowlist allows all", func(t *testing.T) { - cfg := config.WeComConfig{} - cfg.SetToken("test_token") - cfg.WebhookURL = "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=test" - cfg.AllowFrom = []string{} - ch, _ := NewWeComBotChannel(cfg, msgBus) - if !ch.IsAllowed("any_user") { - t.Error("empty allowlist should allow all users") - } - }) - - t.Run("allowlist restricts users", func(t *testing.T) { - cfg := config.WeComConfig{} - cfg.SetToken("test_token") - cfg.WebhookURL = "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=test" - cfg.AllowFrom = []string{"allowed_user"} - ch, _ := NewWeComBotChannel(cfg, msgBus) - if !ch.IsAllowed("allowed_user") { - t.Error("allowed user should pass allowlist check") - } - if ch.IsAllowed("blocked_user") { - t.Error("non-allowed user should be blocked") - } - }) -} - -func TestWeComBotVerifySignature(t *testing.T) { - msgBus := bus.NewMessageBus() - cfg := config.WeComConfig{} - cfg.SetToken("test_token") - cfg.WebhookURL = "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=test" - ch, _ := NewWeComBotChannel(cfg, msgBus) - - t.Run("valid signature", func(t *testing.T) { - timestamp := "1234567890" - nonce := "test_nonce" - msgEncrypt := "test_message" - expectedSig := generateSignature("test_token", timestamp, nonce, msgEncrypt) - - if !verifySignature(ch.config.Token(), expectedSig, timestamp, nonce, msgEncrypt) { - t.Error("valid signature should pass verification") - } - }) - - t.Run("invalid signature", func(t *testing.T) { - timestamp := "1234567890" - nonce := "test_nonce" - msgEncrypt := "test_message" - - if verifySignature(ch.config.Token(), "invalid_sig", timestamp, nonce, msgEncrypt) { - t.Error("invalid signature should fail verification") - } - }) - - t.Run("empty token rejects verification (fail-closed)", func(t *testing.T) { - cfgEmpty := config.WeComConfig{} - cfgEmpty.SetToken("") - cfgEmpty.WebhookURL = "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=test" - chEmpty := &WeComBotChannel{ - config: cfgEmpty, - } - - if verifySignature(chEmpty.config.Token(), "any_sig", "any_ts", "any_nonce", "any_msg") { - t.Error("empty token should reject verification (fail-closed)") - } - }) -} - -func TestWeComBotDecryptMessage(t *testing.T) { - msgBus := bus.NewMessageBus() - - t.Run("decrypt without AES key", func(t *testing.T) { - cfg := config.WeComConfig{} - cfg.SetToken("test_token") - cfg.WebhookURL = "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=test" - cfg.SetEncodingAESKey("") - ch, _ := NewWeComBotChannel(cfg, msgBus) - - // Without AES key, message should be base64 decoded only - plainText := "hello world" - encoded := base64.StdEncoding.EncodeToString([]byte(plainText)) - - result, err := decryptMessage(encoded, ch.config.EncodingAESKey()) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - if result != plainText { - t.Errorf("decryptMessage() = %q, want %q", result, plainText) - } - }) - - t.Run("decrypt with AES key", func(t *testing.T) { - aesKey := generateTestAESKey() - cfg := config.WeComConfig{} - cfg.SetToken("test_token") - cfg.WebhookURL = "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=test" - cfg.SetEncodingAESKey(aesKey) - ch, _ := NewWeComBotChannel(cfg, msgBus) - - originalMsg := "Hello" - encrypted, err := encryptTestMessage(originalMsg, aesKey) - if err != nil { - t.Fatalf("failed to encrypt test message: %v", err) - } - - result, err := decryptMessage(encrypted, ch.config.EncodingAESKey()) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - if result != originalMsg { - t.Errorf("WeComDecryptMessage() = %q, want %q", result, originalMsg) - } - }) - - t.Run("invalid base64", func(t *testing.T) { - cfg := config.WeComConfig{} - cfg.SetToken("test_token") - cfg.WebhookURL = "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=test" - cfg.SetEncodingAESKey("") - ch, _ := NewWeComBotChannel(cfg, msgBus) - - _, err := decryptMessage("invalid_base64!!!", ch.config.EncodingAESKey()) - if err == nil { - t.Error("expected error for invalid base64, got nil") - } - }) - - t.Run("invalid AES key", func(t *testing.T) { - cfg := config.WeComConfig{} - cfg.SetToken("test_token") - cfg.WebhookURL = "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=test" - cfg.SetEncodingAESKey("invalid_key") - ch, _ := NewWeComBotChannel(cfg, msgBus) - - _, err := decryptMessage(base64.StdEncoding.EncodeToString([]byte("test")), ch.config.EncodingAESKey()) - if err == nil { - t.Error("expected error for invalid AES key, got nil") - } - }) -} - -func TestWeComBotPKCS7Unpad(t *testing.T) { - tests := []struct { - name string - input []byte - expected []byte - }{ - { - name: "empty input", - input: []byte{}, - expected: []byte{}, - }, - { - name: "valid padding 3 bytes", - input: append([]byte("hello"), bytes.Repeat([]byte{3}, 3)...), - expected: []byte("hello"), - }, - { - name: "valid padding 16 bytes (full block)", - input: append([]byte("123456789012345"), bytes.Repeat([]byte{16}, 16)...), - expected: []byte("123456789012345"), - }, - { - name: "invalid padding larger than data", - input: []byte{20}, - expected: nil, // should return error - }, - { - name: "invalid padding zero", - input: append([]byte("test"), byte(0)), - expected: nil, // should return error - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - result, err := pkcs7Unpad(tt.input) - if tt.expected == nil { - // This case should return an error - if err == nil { - t.Errorf("pkcs7Unpad() expected error for invalid padding, got result: %v", result) - } - return - } - if err != nil { - t.Errorf("pkcs7Unpad() unexpected error: %v", err) - return - } - if !bytes.Equal(result, tt.expected) { - t.Errorf("pkcs7Unpad() = %v, want %v", result, tt.expected) - } - }) - } -} - -func TestWeComBotHandleVerification(t *testing.T) { - msgBus := bus.NewMessageBus() - aesKey := generateTestAESKey() - cfg := config.WeComConfig{} - cfg.SetToken("test_token") - cfg.SetEncodingAESKey(aesKey) - cfg.WebhookURL = "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=test" - ch, _ := NewWeComBotChannel(cfg, msgBus) - - t.Run("valid verification request", func(t *testing.T) { - echostr := "test_echostr_123" - encryptedEchostr, _ := encryptTestMessage(echostr, aesKey) - timestamp := "1234567890" - nonce := "test_nonce" - signature := generateSignature("test_token", timestamp, nonce, encryptedEchostr) - - req := httptest.NewRequest( - http.MethodGet, - "/webhook/wecom?msg_signature="+signature+"×tamp="+timestamp+"&nonce="+nonce+"&echostr="+encryptedEchostr, - nil, - ) - w := httptest.NewRecorder() - - ch.handleVerification(context.Background(), w, req) - - if w.Code != http.StatusOK { - t.Errorf("status code = %d, want %d", w.Code, http.StatusOK) - } - if w.Body.String() != echostr { - t.Errorf("response body = %q, want %q", w.Body.String(), echostr) - } - }) - - t.Run("missing parameters", func(t *testing.T) { - req := httptest.NewRequest(http.MethodGet, "/webhook/wecom?msg_signature=sig×tamp=ts", nil) - w := httptest.NewRecorder() - - ch.handleVerification(context.Background(), w, req) - - if w.Code != http.StatusBadRequest { - t.Errorf("status code = %d, want %d", w.Code, http.StatusBadRequest) - } - }) - - t.Run("invalid signature", func(t *testing.T) { - echostr := "test_echostr" - encryptedEchostr, _ := encryptTestMessage(echostr, aesKey) - timestamp := "1234567890" - nonce := "test_nonce" - - req := httptest.NewRequest( - http.MethodGet, - "/webhook/wecom?msg_signature=invalid_sig×tamp="+timestamp+"&nonce="+nonce+"&echostr="+encryptedEchostr, - nil, - ) - w := httptest.NewRecorder() - - ch.handleVerification(context.Background(), w, req) - - if w.Code != http.StatusForbidden { - t.Errorf("status code = %d, want %d", w.Code, http.StatusForbidden) - } - }) -} - -func TestWeComBotHandleMessageCallback(t *testing.T) { - msgBus := bus.NewMessageBus() - aesKey := generateTestAESKey() - cfg := config.WeComConfig{} - cfg.SetToken("test_token") - cfg.SetEncodingAESKey(aesKey) - cfg.WebhookURL = "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=test" - ch, _ := NewWeComBotChannel(cfg, msgBus) - - runBotMessageCallback := func(t *testing.T, jsonMsg string) *httptest.ResponseRecorder { - t.Helper() - encrypted, _ := encryptTestMessage(jsonMsg, aesKey) - encryptedWrapper := struct { - XMLName xml.Name `xml:"xml"` - Encrypt string `xml:"Encrypt"` - }{ - Encrypt: encrypted, - } - wrapperData, _ := xml.Marshal(encryptedWrapper) - timestamp := "1234567890" - nonce := "test_nonce" - signature := generateSignature("test_token", timestamp, nonce, encrypted) - req := httptest.NewRequest( - http.MethodPost, - "/webhook/wecom?msg_signature="+signature+"×tamp="+timestamp+"&nonce="+nonce, - bytes.NewReader(wrapperData), - ) - w := httptest.NewRecorder() - ch.handleMessageCallback(context.Background(), w, req) - return w - } - - t.Run("valid direct message callback", func(t *testing.T) { - w := runBotMessageCallback(t, `{ - "msgid": "test_msg_id_123", - "aibotid": "test_aibot_id", - "chattype": "single", - "from": {"userid": "user123"}, - "response_url": "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=test", - "msgtype": "text", - "text": {"content": "Hello World"} - }`) - if w.Code != http.StatusOK { - t.Errorf("status code = %d, want %d", w.Code, http.StatusOK) - } - if w.Body.String() != "success" { - t.Errorf("response body = %q, want %q", w.Body.String(), "success") - } - }) - - t.Run("valid group message callback", func(t *testing.T) { - w := runBotMessageCallback(t, `{ - "msgid": "test_msg_id_456", - "aibotid": "test_aibot_id", - "chatid": "group_chat_id_123", - "chattype": "group", - "from": {"userid": "user456"}, - "response_url": "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=test", - "msgtype": "text", - "text": {"content": "Hello Group"} - }`) - if w.Code != http.StatusOK { - t.Errorf("status code = %d, want %d", w.Code, http.StatusOK) - } - if w.Body.String() != "success" { - t.Errorf("response body = %q, want %q", w.Body.String(), "success") - } - }) - - t.Run("missing parameters", func(t *testing.T) { - req := httptest.NewRequest(http.MethodPost, "/webhook/wecom?msg_signature=sig", nil) - w := httptest.NewRecorder() - - ch.handleMessageCallback(context.Background(), w, req) - - if w.Code != http.StatusBadRequest { - t.Errorf("status code = %d, want %d", w.Code, http.StatusBadRequest) - } - }) - - t.Run("invalid XML", func(t *testing.T) { - timestamp := "1234567890" - nonce := "test_nonce" - signature := generateSignature("test_token", timestamp, nonce, "") - - req := httptest.NewRequest( - http.MethodPost, - "/webhook/wecom?msg_signature="+signature+"×tamp="+timestamp+"&nonce="+nonce, - strings.NewReader("invalid xml"), - ) - w := httptest.NewRecorder() - - ch.handleMessageCallback(context.Background(), w, req) - - if w.Code != http.StatusBadRequest { - t.Errorf("status code = %d, want %d", w.Code, http.StatusBadRequest) - } - }) - - t.Run("invalid signature", func(t *testing.T) { - encryptedWrapper := struct { - XMLName xml.Name `xml:"xml"` - Encrypt string `xml:"Encrypt"` - }{ - Encrypt: "encrypted_data", - } - wrapperData, _ := xml.Marshal(encryptedWrapper) - - timestamp := "1234567890" - nonce := "test_nonce" - - req := httptest.NewRequest( - http.MethodPost, - "/webhook/wecom?msg_signature=invalid_sig×tamp="+timestamp+"&nonce="+nonce, - bytes.NewReader(wrapperData), - ) - w := httptest.NewRecorder() - - ch.handleMessageCallback(context.Background(), w, req) - - if w.Code != http.StatusForbidden { - t.Errorf("status code = %d, want %d", w.Code, http.StatusForbidden) - } - }) -} - -func TestWeComBotProcessMessage(t *testing.T) { - msgBus := bus.NewMessageBus() - cfg := config.WeComConfig{} - cfg.SetToken("test_token") - cfg.WebhookURL = "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=test" - ch, _ := NewWeComBotChannel(cfg, msgBus) - - t.Run("process direct text message", func(t *testing.T) { - msg := WeComBotMessage{ - MsgID: "test_msg_id_123", - AIBotID: "test_aibot_id", - ChatType: "single", - ResponseURL: "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=test", - MsgType: "text", - } - msg.From.UserID = "user123" - msg.Text.Content = "Hello World" - - // Should not panic - ch.processMessage(context.Background(), msg) - }) - - t.Run("process group text message", func(t *testing.T) { - msg := WeComBotMessage{ - MsgID: "test_msg_id_456", - AIBotID: "test_aibot_id", - ChatID: "group_chat_id_123", - ChatType: "group", - ResponseURL: "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=test", - MsgType: "text", - } - msg.From.UserID = "user456" - msg.Text.Content = "Hello Group" - - // Should not panic - ch.processMessage(context.Background(), msg) - }) - - t.Run("process voice message", func(t *testing.T) { - msg := WeComBotMessage{ - MsgID: "test_msg_id_789", - AIBotID: "test_aibot_id", - ChatType: "single", - ResponseURL: "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=test", - MsgType: "voice", - } - msg.From.UserID = "user123" - msg.Voice.Content = "Voice message text" - - // Should not panic - ch.processMessage(context.Background(), msg) - }) - - t.Run("skip unsupported message type", func(t *testing.T) { - msg := WeComBotMessage{ - MsgID: "test_msg_id_000", - AIBotID: "test_aibot_id", - ChatType: "single", - ResponseURL: "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=test", - MsgType: "video", - } - msg.From.UserID = "user123" - - // Should not panic - ch.processMessage(context.Background(), msg) - }) -} - -func TestWeComBotHandleWebhook(t *testing.T) { - msgBus := bus.NewMessageBus() - cfg := config.WeComConfig{} - cfg.SetToken("test_token") - cfg.WebhookURL = "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=test" - ch, _ := NewWeComBotChannel(cfg, msgBus) - - t.Run("GET request calls verification", func(t *testing.T) { - echostr := "test_echostr" - encoded := base64.StdEncoding.EncodeToString([]byte(echostr)) - timestamp := "1234567890" - nonce := "test_nonce" - signature := generateSignature("test_token", timestamp, nonce, encoded) - - req := httptest.NewRequest( - http.MethodGet, - "/webhook/wecom?msg_signature="+signature+"×tamp="+timestamp+"&nonce="+nonce+"&echostr="+encoded, - nil, - ) - w := httptest.NewRecorder() - - ch.handleWebhook(w, req) - - if w.Code != http.StatusOK { - t.Errorf("status code = %d, want %d", w.Code, http.StatusOK) - } - }) - - t.Run("POST request calls message callback", func(t *testing.T) { - encryptedWrapper := struct { - XMLName xml.Name `xml:"xml"` - Encrypt string `xml:"Encrypt"` - }{ - Encrypt: base64.StdEncoding.EncodeToString([]byte("test")), - } - wrapperData, _ := xml.Marshal(encryptedWrapper) - - timestamp := "1234567890" - nonce := "test_nonce" - signature := generateSignature("test_token", timestamp, nonce, encryptedWrapper.Encrypt) - - req := httptest.NewRequest( - http.MethodPost, - "/webhook/wecom?msg_signature="+signature+"×tamp="+timestamp+"&nonce="+nonce, - bytes.NewReader(wrapperData), - ) - w := httptest.NewRecorder() - - ch.handleWebhook(w, req) - - // Should not be method not allowed - if w.Code == http.StatusMethodNotAllowed { - t.Error("POST request should not return Method Not Allowed") - } - }) - - t.Run("unsupported method", func(t *testing.T) { - req := httptest.NewRequest(http.MethodPut, "/webhook/wecom", nil) - w := httptest.NewRecorder() - - ch.handleWebhook(w, req) - - if w.Code != http.StatusMethodNotAllowed { - t.Errorf("status code = %d, want %d", w.Code, http.StatusMethodNotAllowed) - } - }) -} - -func TestWeComBotHandleHealth(t *testing.T) { - msgBus := bus.NewMessageBus() - cfg := config.WeComConfig{} - cfg.SetToken("test_token") - cfg.WebhookURL = "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=test" - ch, _ := NewWeComBotChannel(cfg, msgBus) - - req := httptest.NewRequest(http.MethodGet, "/health/wecom", nil) - w := httptest.NewRecorder() - - ch.handleHealth(w, req) - - if w.Code != http.StatusOK { - t.Errorf("status code = %d, want %d", w.Code, http.StatusOK) - } - - contentType := w.Header().Get("Content-Type") - if contentType != "application/json" { - t.Errorf("Content-Type = %q, want %q", contentType, "application/json") - } - - body := w.Body.String() - if !strings.Contains(body, "status") || !strings.Contains(body, "running") { - t.Errorf("response body should contain status and running fields, got: %s", body) - } -} - -func TestWeComBotReplyMessage(t *testing.T) { - msg := WeComBotReplyMessage{ - MsgType: "text", - } - msg.Text.Content = "Hello World" - - if msg.MsgType != "text" { - t.Errorf("MsgType = %q, want %q", msg.MsgType, "text") - } - if msg.Text.Content != "Hello World" { - t.Errorf("Text.Content = %q, want %q", msg.Text.Content, "Hello World") - } -} - -func TestWeComBotMessageStructure(t *testing.T) { - jsonData := `{ - "msgid": "test_msg_id_123", - "aibotid": "test_aibot_id", - "chatid": "group_chat_id_123", - "chattype": "group", - "from": {"userid": "user123"}, - "response_url": "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=test", - "msgtype": "text", - "text": {"content": "Hello World"} - }` - - var msg WeComBotMessage - err := json.Unmarshal([]byte(jsonData), &msg) - if err != nil { - t.Fatalf("failed to unmarshal JSON: %v", err) - } - - if msg.MsgID != "test_msg_id_123" { - t.Errorf("MsgID = %q, want %q", msg.MsgID, "test_msg_id_123") - } - if msg.AIBotID != "test_aibot_id" { - t.Errorf("AIBotID = %q, want %q", msg.AIBotID, "test_aibot_id") - } - if msg.ChatID != "group_chat_id_123" { - t.Errorf("ChatID = %q, want %q", msg.ChatID, "group_chat_id_123") - } - if msg.ChatType != "group" { - t.Errorf("ChatType = %q, want %q", msg.ChatType, "group") - } - if msg.From.UserID != "user123" { - t.Errorf("From.UserID = %q, want %q", msg.From.UserID, "user123") - } - if msg.MsgType != "text" { - t.Errorf("MsgType = %q, want %q", msg.MsgType, "text") - } - if msg.Text.Content != "Hello World" { - t.Errorf("Text.Content = %q, want %q", msg.Text.Content, "Hello World") - } -} diff --git a/pkg/channels/wecom/common.go b/pkg/channels/wecom/common.go deleted file mode 100644 index 9a622a2fc..000000000 --- a/pkg/channels/wecom/common.go +++ /dev/null @@ -1,199 +0,0 @@ -package wecom - -import ( - "bytes" - "crypto/aes" - "crypto/cipher" - "crypto/rand" - "crypto/sha1" - "encoding/base64" - "encoding/binary" - "fmt" - "math/big" - "sort" - "strings" -) - -// blockSize is the PKCS7 block size used by WeCom (32) -const blockSize = 32 - -// computeSignature computes the WeCom message signature from the given parameters. -// It sorts [token, timestamp, nonce, encrypt], concatenates them and returns the SHA1 hex digest. -func computeSignature(token, timestamp, nonce, encrypt string) string { - params := []string{token, timestamp, nonce, encrypt} - sort.Strings(params) - str := strings.Join(params, "") - hash := sha1.Sum([]byte(str)) - return fmt.Sprintf("%x", hash) -} - -// verifySignature verifies the message signature for WeCom -// This is a common function used by both WeCom Bot and WeCom App -func verifySignature(token, msgSignature, timestamp, nonce, msgEncrypt string) bool { - if token == "" { - return false - } - return computeSignature(token, timestamp, nonce, msgEncrypt) == msgSignature -} - -// decryptMessage decrypts the encrypted message using AES -// For AIBOT, receiveid should be the aibotid; for other apps, it should be corp_id -func decryptMessage(encryptedMsg, encodingAESKey string) (string, error) { - return decryptMessageWithVerify(encryptedMsg, encodingAESKey, "") -} - -// decryptMessageWithVerify decrypts the encrypted message and optionally verifies receiveid -// receiveid: for AIBOT use aibotid, for WeCom App use corp_id. If empty, skip verification. -func decryptMessageWithVerify(encryptedMsg, encodingAESKey, receiveid string) (string, error) { - if encodingAESKey == "" { - // No encryption, return as is (base64 decode) - decoded, err := base64.StdEncoding.DecodeString(encryptedMsg) - if err != nil { - return "", err - } - return string(decoded), nil - } - - aesKey, err := decodeWeComAESKey(encodingAESKey) - if err != nil { - return "", err - } - - cipherText, err := base64.StdEncoding.DecodeString(encryptedMsg) - if err != nil { - return "", fmt.Errorf("failed to decode message: %w", err) - } - - plainText, err := decryptAESCBC(aesKey, cipherText) - if err != nil { - return "", err - } - - return unpackWeComFrame(plainText, receiveid) -} - -// decodeWeComAESKey base64-decodes the 43-character EncodingAESKey (trailing "=" is -// appended automatically) and validates that the result is exactly 32 bytes. -// It is the single place that handles this repeated pattern in both encrypt and decrypt paths. -func decodeWeComAESKey(encodingAESKey string) ([]byte, error) { - aesKey, err := base64.StdEncoding.DecodeString(encodingAESKey + "=") - if err != nil { - return nil, fmt.Errorf("failed to decode AES key: %w", err) - } - if len(aesKey) != 32 { - return nil, fmt.Errorf("invalid AES key length: %d", len(aesKey)) - } - return aesKey, nil -} - -// encryptAESCBC encrypts plaintext using AES-CBC with the given key, mirroring -// decryptAESCBC. IV = aesKey[:aes.BlockSize]. The caller must PKCS7-pad the -// plaintext to a multiple of aes.BlockSize before calling. -func encryptAESCBC(aesKey, plaintext []byte) ([]byte, error) { - block, err := aes.NewCipher(aesKey) - if err != nil { - return nil, fmt.Errorf("failed to create cipher: %w", err) - } - iv := aesKey[:aes.BlockSize] - ciphertext := make([]byte, len(plaintext)) - cipher.NewCBCEncrypter(block, iv).CryptBlocks(ciphertext, plaintext) - return ciphertext, nil -} - -// packWeComFrame builds the WeCom wire format: -// -// random(16 ASCII digits) + msg_len(4, big-endian) + msg + receiveid -func packWeComFrame(msg, receiveid string) ([]byte, error) { - randomBytes := make([]byte, 16) - for i := range 16 { - n, err := rand.Int(rand.Reader, big.NewInt(10)) - if err != nil { - return nil, fmt.Errorf("failed to generate random: %w", err) - } - randomBytes[i] = byte('0' + n.Int64()) - } - msgBytes := []byte(msg) - msgLenBytes := make([]byte, 4) - binary.BigEndian.PutUint32(msgLenBytes, uint32(len(msgBytes))) - var buf bytes.Buffer - buf.Write(randomBytes) - buf.Write(msgLenBytes) - buf.Write(msgBytes) - buf.WriteString(receiveid) - return buf.Bytes(), nil -} - -// unpackWeComFrame parses the WeCom wire format produced by packWeComFrame. -// If receiveid is non-empty it verifies the frame's trailing receiveid field. -func unpackWeComFrame(data []byte, receiveid string) (string, error) { - if len(data) < 20 { - return "", fmt.Errorf("decrypted frame too short: %d bytes", len(data)) - } - msgLen := binary.BigEndian.Uint32(data[16:20]) - if int(msgLen) > len(data)-20 { - return "", fmt.Errorf("invalid message length: %d", msgLen) - } - msg := data[20 : 20+msgLen] - if receiveid != "" && len(data) > 20+int(msgLen) { - actualReceiveID := string(data[20+msgLen:]) - if actualReceiveID != receiveid { - return "", fmt.Errorf("receiveid mismatch: expected %s, got %s", receiveid, actualReceiveID) - } - } - return string(msg), nil -} - -// decryptAESCBC decrypts ciphertext using AES-CBC with the given key. -// IV = aesKey[:aes.BlockSize]. PKCS7 padding is stripped from the returned plaintext. -func decryptAESCBC(aesKey, ciphertext []byte) ([]byte, error) { - if len(ciphertext) == 0 { - return nil, fmt.Errorf("ciphertext is empty") - } - if len(ciphertext)%aes.BlockSize != 0 { - return nil, fmt.Errorf("ciphertext length %d is not a multiple of block size", len(ciphertext)) - } - block, err := aes.NewCipher(aesKey) - if err != nil { - return nil, fmt.Errorf("failed to create cipher: %w", err) - } - iv := aesKey[:aes.BlockSize] - plaintext := make([]byte, len(ciphertext)) - cipher.NewCBCDecrypter(block, iv).CryptBlocks(plaintext, ciphertext) - plaintext, err = pkcs7Unpad(plaintext) - if err != nil { - return nil, fmt.Errorf("failed to unpad: %w", err) - } - return plaintext, nil -} - -// pkcs7Pad adds PKCS7 padding -func pkcs7Pad(data []byte, blockSize int) []byte { - padding := blockSize - (len(data) % blockSize) - if padding == 0 { - padding = blockSize - } - padText := bytes.Repeat([]byte{byte(padding)}, padding) - return append(data, padText...) -} - -// pkcs7Unpad removes PKCS7 padding with validation -func pkcs7Unpad(data []byte) ([]byte, error) { - if len(data) == 0 { - return data, nil - } - padding := int(data[len(data)-1]) - // WeCom uses 32-byte block size for PKCS7 padding - if padding == 0 || padding > blockSize { - return nil, fmt.Errorf("invalid padding size: %d", padding) - } - if padding > len(data) { - return nil, fmt.Errorf("padding size larger than data") - } - // Verify all padding bytes - for i := range padding { - if data[len(data)-1-i] != byte(padding) { - return nil, fmt.Errorf("invalid padding byte at position %d", i) - } - } - return data[:len(data)-padding], nil -} diff --git a/pkg/channels/wecom/dedupe.go b/pkg/channels/wecom/dedupe.go deleted file mode 100644 index 865be668e..000000000 --- a/pkg/channels/wecom/dedupe.go +++ /dev/null @@ -1,54 +0,0 @@ -package wecom - -import "sync" - -const wecomMaxProcessedMessages = 1000 - -// MessageDeduplicator provides thread-safe message deduplication using a circular queue (ring buffer) -// combined with a hash map. This ensures fast O(1) lookups while naturally evicting the oldest -// messages without causing "amnesia cliffs" when the limit is reached. -type MessageDeduplicator struct { - mu sync.Mutex - msgs map[string]bool - ring []string - idx int - max int -} - -// NewMessageDeduplicator creates a new deduplicator with the specified capacity. -func NewMessageDeduplicator(maxEntries int) *MessageDeduplicator { - if maxEntries <= 0 { - maxEntries = wecomMaxProcessedMessages - } - return &MessageDeduplicator{ - msgs: make(map[string]bool, maxEntries), - ring: make([]string, maxEntries), - max: maxEntries, - } -} - -// MarkMessageProcessed marks msgID as processed and returns false for duplicates. -func (d *MessageDeduplicator) MarkMessageProcessed(msgID string) bool { - d.mu.Lock() - defer d.mu.Unlock() - - // 1. Check for duplicate - if d.msgs[msgID] { - return false - } - - // 2. Evict the oldest message at our current ring position (if any) - oldestID := d.ring[d.idx] - if oldestID != "" { - delete(d.msgs, oldestID) - } - - // 3. Store the new message - d.msgs[msgID] = true - d.ring[d.idx] = msgID - - // 4. Advance the circle queue index - d.idx = (d.idx + 1) % d.max - - return true -} diff --git a/pkg/channels/wecom/dedupe_test.go b/pkg/channels/wecom/dedupe_test.go deleted file mode 100644 index 10dff4cfe..000000000 --- a/pkg/channels/wecom/dedupe_test.go +++ /dev/null @@ -1,83 +0,0 @@ -package wecom - -import ( - "sync" - "testing" -) - -func TestMessageDeduplicator_DuplicateDetection(t *testing.T) { - d := NewMessageDeduplicator(wecomMaxProcessedMessages) - - if ok := d.MarkMessageProcessed("msg-1"); !ok { - t.Fatalf("first message should be accepted") - } - - if ok := d.MarkMessageProcessed("msg-1"); ok { - t.Fatalf("duplicate message should be rejected") - } -} - -func TestMessageDeduplicator_ConcurrentSameMessage(t *testing.T) { - d := NewMessageDeduplicator(wecomMaxProcessedMessages) - - const goroutines = 64 - var wg sync.WaitGroup - wg.Add(goroutines) - - results := make(chan bool, goroutines) - for i := 0; i < goroutines; i++ { - go func() { - defer wg.Done() - results <- d.MarkMessageProcessed("msg-concurrent") - }() - } - - wg.Wait() - close(results) - - successes := 0 - for ok := range results { - if ok { - successes++ - } - } - - if successes != 1 { - t.Fatalf("expected exactly 1 successful mark, got %d", successes) - } -} - -func TestMessageDeduplicator_CircularQueueEviction(t *testing.T) { - // Create a deduplicator with a very small capacity to test eviction easily. - capacity := 3 - d := NewMessageDeduplicator(capacity) - - // Fill the queue. - d.MarkMessageProcessed("msg-1") - d.MarkMessageProcessed("msg-2") - d.MarkMessageProcessed("msg-3") - - // At this point, the queue is full. msg-1 is the oldest. - if len(d.msgs) != 3 { - t.Fatalf("expected map size to be 3, got %d", len(d.msgs)) - } - - // This should evict msg-1 and add msg-4. - if ok := d.MarkMessageProcessed("msg-4"); !ok { - t.Fatalf("msg-4 should be accepted") - } - - if len(d.msgs) != 3 { - t.Fatalf("expected map size to remain at max capacity (3), got %d", len(d.msgs)) - } - - // msg-1 should now be forgotten (evicted). - if ok := d.MarkMessageProcessed("msg-1"); !ok { - t.Fatalf("msg-1 should be accepted again because it was evicted") - } - - // msg-2 should have been evicted when we added msg-1 back. - if ok := d.MarkMessageProcessed("msg-2"); !ok { - t.Fatalf("msg-2 should be accepted again because it was evicted") - } -} diff --git a/pkg/channels/wecom/init.go b/pkg/channels/wecom/init.go index bc5a70fa3..3aad84d42 100644 --- a/pkg/channels/wecom/init.go +++ b/pkg/channels/wecom/init.go @@ -8,12 +8,6 @@ import ( func init() { channels.RegisterFactory("wecom", func(cfg *config.Config, b *bus.MessageBus) (channels.Channel, error) { - return NewWeComBotChannel(cfg.Channels.WeCom, b) - }) - channels.RegisterFactory("wecom_app", func(cfg *config.Config, b *bus.MessageBus) (channels.Channel, error) { - return NewWeComAppChannel(cfg.Channels.WeComApp, b) - }) - channels.RegisterFactory("wecom_aibot", func(cfg *config.Config, b *bus.MessageBus) (channels.Channel, error) { - return NewWeComAIBotChannel(cfg.Channels.WeComAIBot, b) + return NewChannel(cfg.Channels.WeCom, b) }) } diff --git a/pkg/channels/wecom/media.go b/pkg/channels/wecom/media.go new file mode 100644 index 000000000..974a3bf4d --- /dev/null +++ b/pkg/channels/wecom/media.go @@ -0,0 +1,802 @@ +package wecom + +import ( + "context" + "crypto/aes" + "crypto/cipher" + "crypto/md5" + "encoding/base64" + "encoding/hex" + "encoding/json" + "fmt" + "io" + "mime" + "net/http" + "net/url" + "os" + "path/filepath" + "strings" + "time" + + "github.com/h2non/filetype" + + "github.com/sipeed/picoclaw/pkg/bus" + "github.com/sipeed/picoclaw/pkg/media" +) + +const ( + wecomOutboundMediaMaxBytes = 20 << 20 + wecomOutboundImageMaxBytes = 2 << 20 + wecomOutboundVoiceMaxBytes = 2 << 20 + wecomOutboundVideoMaxBytes = 10 << 20 + wecomUploadChunkMaxBytes = 512 << 10 + wecomUploadMaxChunks = 100 + wecomUploadMinBytes = 5 +) + +type wecomOutboundMedia struct { + MsgType string + MediaID string + Title string + Description string +} + +func (m *wecomOutboundMedia) respondBody() wecomRespondMsgBody { + body := wecomRespondMsgBody{MsgType: m.MsgType} + switch m.MsgType { + case "file": + body.File = &wecomMediaRefContent{MediaID: m.MediaID} + case "image": + body.Image = &wecomMediaRefContent{MediaID: m.MediaID} + case "voice": + body.Voice = &wecomMediaRefContent{MediaID: m.MediaID} + case "video": + body.Video = &wecomVideoContent{ + MediaID: m.MediaID, + Title: m.Title, + Description: m.Description, + } + } + return body +} + +func (m *wecomOutboundMedia) sendBody(chatID string, chatType uint32) wecomSendMsgBody { + body := wecomSendMsgBody{ + ChatID: chatID, + ChatType: chatType, + MsgType: m.MsgType, + } + switch m.MsgType { + case "file": + body.File = &wecomMediaRefContent{MediaID: m.MediaID} + case "image": + body.Image = &wecomMediaRefContent{MediaID: m.MediaID} + case "voice": + body.Voice = &wecomMediaRefContent{MediaID: m.MediaID} + case "video": + body.Video = &wecomVideoContent{ + MediaID: m.MediaID, + Title: m.Title, + Description: m.Description, + } + } + return body +} + +func decodeMediaAESKey(value string) ([]byte, error) { + if value == "" { + return nil, nil + } + key, err := base64.StdEncoding.DecodeString(value) + if err == nil && len(key) == 32 { + return key, nil + } + key, err = base64.StdEncoding.DecodeString(value + "=") + if err != nil { + return nil, fmt.Errorf("decode AES key: %w", err) + } + if len(key) != 32 { + return nil, fmt.Errorf("invalid AES key length %d", len(key)) + } + return key, nil +} + +func decryptAESCBC(key, ciphertext []byte) ([]byte, error) { + if len(ciphertext) == 0 { + return nil, fmt.Errorf("ciphertext is empty") + } + if len(ciphertext)%aes.BlockSize != 0 { + return nil, fmt.Errorf("ciphertext length %d is not a multiple of block size", len(ciphertext)) + } + block, err := aes.NewCipher(key) + if err != nil { + return nil, fmt.Errorf("create cipher: %w", err) + } + plaintext := make([]byte, len(ciphertext)) + iv := key[:aes.BlockSize] + cipher.NewCBCDecrypter(block, iv).CryptBlocks(plaintext, ciphertext) + return pkcs7Unpad(plaintext) +} + +func pkcs7Unpad(data []byte) ([]byte, error) { + if len(data) == 0 { + return nil, fmt.Errorf("empty plaintext") + } + padding := int(data[len(data)-1]) + if padding == 0 || padding > 32 || padding > len(data) { + return nil, fmt.Errorf("invalid padding size %d", padding) + } + for i := 0; i < padding; i++ { + if data[len(data)-1-i] != byte(padding) { + return nil, fmt.Errorf("invalid padding byte") + } + } + return data[:len(data)-padding], nil +} + +func inferMediaExt(contentType, fallback string) string { + contentType = normalizeWeComContentType(contentType) + switch contentType { + case "image/jpeg", "image/jpg": + return ".jpg" + case "image/png": + return ".png" + case "image/gif": + return ".gif" + case "image/webp": + return ".webp" + case "application/pdf": + return ".pdf" + case "video/mp4": + return ".mp4" + default: + return fallback + } +} + +func normalizeWeComContentType(value string) string { + value = strings.ToLower(strings.TrimSpace(value)) + if idx := strings.Index(value, ";"); idx >= 0 { + value = strings.TrimSpace(value[:idx]) + } + return value +} + +func isGenericWeComContentType(value string) bool { + switch normalizeWeComContentType(value) { + case "", "application/octet-stream", "binary/octet-stream", "application/unknown", "application/binary": + return true + default: + return false + } +} + +func sanitizeWeComFilename(name string) string { + name = filepath.Base(strings.TrimSpace(name)) + if name == "." || name == "/" || name == "" { + return "" + } + return name +} + +func candidateWeComFilename(resourceURL, contentDisposition, fallbackName string) string { + if _, params, err := mime.ParseMediaType(contentDisposition); err == nil { + if name := sanitizeWeComFilename(params["filename"]); name != "" { + return name + } + if name := sanitizeWeComFilename(params["filename*"]); name != "" { + return name + } + } + + if parsed, err := url.Parse(resourceURL); err == nil { + query := parsed.Query() + for _, key := range []string{"filename", "file_name", "name"} { + if name := sanitizeWeComFilename(query.Get(key)); name != "" { + return name + } + } + if name := sanitizeWeComFilename(parsed.Path); name != "" { + return name + } + } + + return sanitizeWeComFilename(fallbackName) +} + +func detectWeComFiletype(data []byte) (string, string) { + kind, err := filetype.Match(data) + if err != nil || kind == filetype.Unknown { + return "", "" + } + ext := "" + if kind.Extension != "" { + ext = "." + strings.ToLower(kind.Extension) + } + return normalizeWeComContentType(kind.MIME.Value), ext +} + +func detectWeComMediaMetadata( + data []byte, + fallbackName, fallbackContentType, resourceURL, contentDisposition string, +) (string, string) { + filename := candidateWeComFilename(resourceURL, contentDisposition, fallbackName) + if filename == "" { + filename = "media" + } + + ext := strings.ToLower(filepath.Ext(filename)) + contentType := normalizeWeComContentType(fallbackContentType) + detectedType, detectedExt := detectWeComFiletype(data) + + if ext != "" && isGenericWeComContentType(contentType) { + if byExt := normalizeWeComContentType(mime.TypeByExtension(ext)); byExt != "" { + contentType = byExt + } + } + + if detectedType != "" { + switch { + case contentType == "": + contentType = detectedType + case isGenericWeComContentType(contentType): + contentType = detectedType + case strings.HasPrefix(detectedType, "image/") && !strings.HasPrefix(contentType, "image/"): + contentType = detectedType + case strings.HasPrefix(detectedType, "audio/") && !strings.HasPrefix(contentType, "audio/"): + contentType = detectedType + case strings.HasPrefix(detectedType, "video/") && !strings.HasPrefix(contentType, "video/"): + contentType = detectedType + } + } + + if contentType == "" && ext != "" { + contentType = normalizeWeComContentType(mime.TypeByExtension(ext)) + } + if contentType == "" { + contentType = normalizeWeComContentType(http.DetectContentType(data)) + } + + if ext == "" { + ext = detectedExt + } + if ext == "" && contentType != "" { + if exts, err := mime.ExtensionsByType(contentType); err == nil && len(exts) > 0 { + ext = strings.ToLower(exts[0]) + } + } + + if filepath.Ext(filename) == "" && ext != "" { + filename += ext + } + return filename, contentType +} + +func (c *WeComChannel) storeRemoteMedia( + ctx context.Context, + scope, msgID, resourceURL, aesKey, fallbackExt string, +) (string, error) { + store := c.GetMediaStore() + if store == nil { + return "", fmt.Errorf("no media store available") + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, resourceURL, nil) + if err != nil { + return "", fmt.Errorf("create request: %w", err) + } + resp, err := c.mediaClient.Do(req) + if err != nil { + return "", fmt.Errorf("download media: %w", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return "", fmt.Errorf("download media returned HTTP %d", resp.StatusCode) + } + + data, err := io.ReadAll(io.LimitReader(resp.Body, wecomOutboundMediaMaxBytes+1)) + if err != nil { + return "", fmt.Errorf("read media: %w", err) + } + if len(data) > wecomOutboundMediaMaxBytes { + return "", fmt.Errorf("media too large") + } + + if aesKey != "" { + key, keyErr := decodeMediaAESKey(aesKey) + if keyErr != nil { + return "", keyErr + } + data, err = decryptAESCBC(key, data) + if err != nil { + return "", fmt.Errorf("decrypt media: %w", err) + } + } + + filename, contentType := detectWeComMediaMetadata( + data, + msgID+fallbackExt, + resp.Header.Get("Content-Type"), + resourceURL, + resp.Header.Get("Content-Disposition"), + ) + ext := filepath.Ext(filename) + if ext == "" { + ext = inferMediaExt(contentType, fallbackExt) + } + mediaDir := filepath.Join(os.TempDir(), "picoclaw_media") + if mkdirErr := os.MkdirAll(mediaDir, 0o700); mkdirErr != nil { + return "", fmt.Errorf("mkdir media dir: %w", mkdirErr) + } + tmpFile, err := os.CreateTemp(mediaDir, msgID+"-*"+ext) + if err != nil { + return "", fmt.Errorf("create temp file: %w", err) + } + tmpPath := tmpFile.Name() + if _, writeErr := tmpFile.Write(data); writeErr != nil { + tmpFile.Close() + _ = os.Remove(tmpPath) + return "", fmt.Errorf("write temp file: %w", writeErr) + } + if closeErr := tmpFile.Close(); closeErr != nil { + _ = os.Remove(tmpPath) + return "", fmt.Errorf("close temp file: %w", closeErr) + } + + ref, err := store.Store(tmpPath, media.MediaMeta{ + Filename: filename, + ContentType: contentType, + Source: "wecom", + CleanupPolicy: media.CleanupPolicyDeleteOnCleanup, + }, scope) + if err != nil { + _ = os.Remove(tmpPath) + return "", err + } + return ref, nil +} + +func detectLocalWeComContentType(localPath, hint string) string { + contentType := normalizeWeComContentType(hint) + if !isGenericWeComContentType(contentType) { + return contentType + } + + if kind, err := filetype.MatchFile(localPath); err == nil && kind != filetype.Unknown { + return normalizeWeComContentType(kind.MIME.Value) + } + + if ext := strings.ToLower(filepath.Ext(localPath)); ext != "" { + if byExt := normalizeWeComContentType(mime.TypeByExtension(ext)); byExt != "" { + return byExt + } + } + + file, err := os.Open(localPath) + if err != nil { + return contentType + } + defer file.Close() + + buf := make([]byte, 512) + n, err := file.Read(buf) + if err != nil && err != io.EOF { + return contentType + } + if n == 0 { + return contentType + } + return normalizeWeComContentType(http.DetectContentType(buf[:n])) +} + +func writeWeComTempFile(prefix, filename string, data []byte) (string, error) { + mediaDir := filepath.Join(os.TempDir(), "picoclaw_media") + if err := os.MkdirAll(mediaDir, 0o700); err != nil { + return "", fmt.Errorf("mkdir media dir: %w", err) + } + + ext := strings.ToLower(filepath.Ext(filename)) + tmpFile, err := os.CreateTemp(mediaDir, prefix+"-*"+ext) + if err != nil { + return "", fmt.Errorf("create temp file: %w", err) + } + tmpPath := tmpFile.Name() + + if _, err := tmpFile.Write(data); err != nil { + _ = tmpFile.Close() + _ = os.Remove(tmpPath) + return "", fmt.Errorf("write temp file: %w", err) + } + if err := tmpFile.Close(); err != nil { + _ = os.Remove(tmpPath) + return "", fmt.Errorf("close temp file: %w", err) + } + return tmpPath, nil +} + +func (c *WeComChannel) downloadRemoteMediaToTemp( + ctx context.Context, + resourceURL, fallbackName string, +) (string, string, string, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, resourceURL, nil) + if err != nil { + return "", "", "", fmt.Errorf("create request: %w", err) + } + + resp, err := c.mediaClient.Do(req) + if err != nil { + return "", "", "", fmt.Errorf("download media: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(io.LimitReader(resp.Body, 1024)) + return "", "", "", fmt.Errorf("download media returned HTTP %d: %s", resp.StatusCode, string(body)) + } + + data, err := io.ReadAll(io.LimitReader(resp.Body, wecomOutboundMediaMaxBytes+1)) + if err != nil { + return "", "", "", fmt.Errorf("read media: %w", err) + } + if len(data) > wecomOutboundMediaMaxBytes { + return "", "", "", fmt.Errorf("media too large") + } + + filename, contentType := detectWeComMediaMetadata( + data, + fallbackName, + resp.Header.Get("Content-Type"), + resourceURL, + resp.Header.Get("Content-Disposition"), + ) + tmpPath, err := writeWeComTempFile("wecom-outbound", filename, data) + if err != nil { + return "", "", "", err + } + return tmpPath, filename, contentType, nil +} + +func (c *WeComChannel) resolveOutboundPart( + ctx context.Context, + part bus.MediaPart, +) (string, string, string, func(), error) { + cleanup := func() {} + filename := sanitizeWeComFilename(part.Filename) + contentType := normalizeWeComContentType(part.ContentType) + ref := strings.TrimSpace(part.Ref) + + switch { + case ref == "": + return "", filename, contentType, cleanup, nil + + case strings.HasPrefix(ref, "http://") || strings.HasPrefix(ref, "https://"): + localPath, name, ct, err := c.downloadRemoteMediaToTemp(ctx, ref, filename) + if err != nil { + return "", "", "", cleanup, err + } + return localPath, name, ct, func() { _ = os.Remove(localPath) }, nil + + case strings.HasPrefix(ref, "media://"): + store := c.GetMediaStore() + if store == nil { + return "", "", "", cleanup, fmt.Errorf("no media store available") + } + + localPath, meta, err := store.ResolveWithMeta(ref) + if err != nil { + return "", "", "", cleanup, err + } + if filename == "" { + filename = sanitizeWeComFilename(meta.Filename) + } + if contentType == "" { + contentType = normalizeWeComContentType(meta.ContentType) + } + if strings.HasPrefix(localPath, "http://") || strings.HasPrefix(localPath, "https://") { + tmpPath, name, ct, err := c.downloadRemoteMediaToTemp(ctx, localPath, filename) + if err != nil { + return "", "", "", cleanup, err + } + return tmpPath, name, ct, func() { _ = os.Remove(tmpPath) }, nil + } + if _, err := os.Stat(localPath); err != nil { + return "", "", "", cleanup, err + } + if filename == "" { + filename = sanitizeWeComFilename(filepath.Base(localPath)) + } + if contentType == "" { + contentType = detectLocalWeComContentType(localPath, "") + } + return localPath, filename, contentType, cleanup, nil + + case strings.HasPrefix(ref, "file://"): + u, err := url.Parse(ref) + if err != nil { + return "", "", "", cleanup, err + } + localPath := u.Path + if _, err := os.Stat(localPath); err != nil { + return "", "", "", cleanup, err + } + if filename == "" { + filename = sanitizeWeComFilename(filepath.Base(localPath)) + } + if contentType == "" { + contentType = detectLocalWeComContentType(localPath, "") + } + return localPath, filename, contentType, cleanup, nil + + default: + if _, err := os.Stat(ref); err != nil { + return "", "", "", cleanup, err + } + if filename == "" { + filename = sanitizeWeComFilename(filepath.Base(ref)) + } + if contentType == "" { + contentType = detectLocalWeComContentType(ref, "") + } + return ref, filename, contentType, cleanup, nil + } +} + +func canWeComSendImage(contentType, ext string, size int64) bool { + if size > wecomOutboundImageMaxBytes { + return false + } + switch normalizeWeComContentType(contentType) { + case "image/jpeg", "image/jpg", "image/png", "image/gif": + return true + } + switch strings.ToLower(ext) { + case ".jpg", ".jpeg", ".png", ".gif": + return true + default: + return false + } +} + +func canWeComSendVoice(contentType, ext string, size int64) bool { + if size > wecomOutboundVoiceMaxBytes { + return false + } + contentType = normalizeWeComContentType(contentType) + return strings.Contains(contentType, "amr") || strings.EqualFold(ext, ".amr") +} + +func canWeComSendVideo(contentType, ext string, size int64) bool { + if size > wecomOutboundVideoMaxBytes { + return false + } + return normalizeWeComContentType(contentType) == "video/mp4" || strings.EqualFold(ext, ".mp4") +} + +func outboundWeComMediaKind(partType, filename, contentType string, size int64) string { + if size < wecomUploadMinBytes { + return "" + } + + partType = strings.ToLower(strings.TrimSpace(partType)) + contentType = normalizeWeComContentType(contentType) + ext := strings.ToLower(filepath.Ext(filename)) + + if partType == "file" { + if size <= wecomOutboundMediaMaxBytes { + return "file" + } + return "" + } + + if (partType == "image" || partType == "") && canWeComSendImage(contentType, ext, size) { + return "image" + } + if (partType == "audio" || partType == "voice" || partType == "") && canWeComSendVoice(contentType, ext, size) { + return "voice" + } + if (partType == "video" || partType == "") && canWeComSendVideo(contentType, ext, size) { + return "video" + } + if size <= wecomOutboundMediaMaxBytes { + return "file" + } + return "" +} + +func trimWeComBytes(value string, limit int) string { + value = strings.TrimSpace(value) + if limit <= 0 || len(value) <= limit { + return value + } + size := 0 + var out strings.Builder + for _, r := range value { + width := len(string(r)) + if size+width > limit { + break + } + size += width + out.WriteRune(r) + } + return out.String() +} + +func ensureWeComOutboundFilename(filename, localPath, contentType string) string { + filename = sanitizeWeComFilename(filename) + if filename == "" { + filename = sanitizeWeComFilename(filepath.Base(localPath)) + } + if filename == "" { + filename = "media" + } + if filepath.Ext(filename) == "" { + fallbackExt := inferMediaExt(contentType, strings.ToLower(filepath.Ext(localPath))) + if fallbackExt != "" { + filename += fallbackExt + } + } + filename = trimWeComBytes(filename, 256) + if filename == "" { + return "media" + } + return filename +} + +func buildWeComVideoContent(mediaID, filename, description string) *wecomVideoContent { + title := strings.TrimSuffix(filename, filepath.Ext(filename)) + title = trimWeComBytes(title, 64) + if title == "" { + title = "video" + } + description = trimWeComBytes(description, 512) + return &wecomVideoContent{ + MediaID: mediaID, + Title: title, + Description: description, + } +} + +func decodeWeComEnvelopeBody[T any](env wecomEnvelope) (T, error) { + var out T + if len(env.Body) == 0 { + return out, fmt.Errorf("wecom response body is empty") + } + if err := json.Unmarshal(env.Body, &out); err != nil { + return out, fmt.Errorf("decode wecom response body: %w", err) + } + return out, nil +} + +func (c *WeComChannel) uploadOutboundMedia( + ctx context.Context, + localPath, filename, contentType string, + part bus.MediaPart, +) (*wecomOutboundMedia, error) { + _ = ctx + + contentType = detectLocalWeComContentType(localPath, contentType) + filename = ensureWeComOutboundFilename(filename, localPath, contentType) + + data, err := os.ReadFile(localPath) + if err != nil { + return nil, fmt.Errorf("read media file: %w", err) + } + size := int64(len(data)) + kind := outboundWeComMediaKind(part.Type, filename, contentType, size) + if kind == "" { + return nil, fmt.Errorf("unsupported wecom media type or size for %q", filename) + } + + totalChunks := (len(data) + wecomUploadChunkMaxBytes - 1) / wecomUploadChunkMaxBytes + if totalChunks <= 0 || totalChunks > wecomUploadMaxChunks { + return nil, fmt.Errorf("wecom upload requires 1-%d chunks, got %d", wecomUploadMaxChunks, totalChunks) + } + + sum := md5.Sum(data) + initEnv, err := c.sendCommandAck(wecomCommand{ + Cmd: wecomCmdUploadMediaInit, + Headers: wecomHeaders{ReqID: randomID(10)}, + Body: wecomUploadMediaInitBody{ + Type: kind, + Filename: filename, + TotalSize: size, + TotalChunks: totalChunks, + MD5: hex.EncodeToString(sum[:]), + }, + }, wecomUploadTimeout) + if err != nil { + return nil, err + } + initResp, err := decodeWeComEnvelopeBody[wecomUploadMediaInitResponse](initEnv) + if err != nil { + return nil, err + } + if strings.TrimSpace(initResp.UploadID) == "" { + return nil, fmt.Errorf("wecom upload init returned empty upload_id") + } + + for idx, offset := 0, 0; offset < len(data); idx, offset = idx+1, offset+wecomUploadChunkMaxBytes { + end := offset + wecomUploadChunkMaxBytes + if end > len(data) { + end = len(data) + } + sendErr := c.sendCommand(wecomCommand{ + Cmd: wecomCmdUploadMediaChunk, + Headers: wecomHeaders{ReqID: randomID(10)}, + Body: wecomUploadMediaChunkBody{ + UploadID: initResp.UploadID, + ChunkIndex: idx, + Base64Data: base64.StdEncoding.EncodeToString(data[offset:end]), + }, + }, wecomUploadTimeout) + if sendErr != nil { + return nil, sendErr + } + } + + finishEnv, err := c.sendCommandAck(wecomCommand{ + Cmd: wecomCmdUploadMediaEnd, + Headers: wecomHeaders{ReqID: randomID(10)}, + Body: wecomUploadMediaFinishBody{ + UploadID: initResp.UploadID, + }, + }, wecomUploadTimeout) + if err != nil { + return nil, err + } + finishResp, err := decodeWeComEnvelopeBody[wecomUploadMediaFinishResponse](finishEnv) + if err != nil { + return nil, err + } + if strings.TrimSpace(finishResp.MediaID) == "" { + return nil, fmt.Errorf("wecom upload finish returned empty media_id") + } + + uploaded := &wecomOutboundMedia{ + MsgType: kind, + MediaID: finishResp.MediaID, + } + if kind == "video" { + video := buildWeComVideoContent(finishResp.MediaID, filename, part.Caption) + uploaded.Title = video.Title + uploaded.Description = video.Description + } + return uploaded, nil +} + +func fallbackWeComMediaText(part bus.MediaPart, kind, filename string) string { + var lines []string + if caption := strings.TrimSpace(part.Caption); caption != "" { + lines = append(lines, caption) + } + + label := kind + if label == "" { + label = "media" + } + if filename != "" { + lines = append(lines, fmt.Sprintf("[%s: %s]", label, filename)) + } else { + lines = append(lines, fmt.Sprintf("[%s attachment]", label)) + } + + ref := strings.TrimSpace(part.Ref) + if strings.HasPrefix(ref, "http://") || strings.HasPrefix(ref, "https://") { + lines = append(lines, ref) + } + + return strings.Join(lines, "\n") +} + +func (c *WeComChannel) resolveMediaRoute(chatID string) (wecomTurn, uint32, bool) { + if turn, ok := c.getTurn(chatID); ok { + if time.Since(turn.CreatedAt) <= wecomStreamMaxDuration { + return turn, turn.ChatType, true + } + c.deleteTurn(chatID) + } + if route, ok := c.routes.Get(chatID); ok { + return wecomTurn{ChatID: route.ChatID, ChatType: route.ChatType}, route.ChatType, false + } + return wecomTurn{ChatID: chatID}, 0, false +} diff --git a/pkg/channels/wecom/media_test.go b/pkg/channels/wecom/media_test.go new file mode 100644 index 000000000..d5307e5d2 --- /dev/null +++ b/pkg/channels/wecom/media_test.go @@ -0,0 +1,180 @@ +package wecom + +import ( + "bytes" + "context" + "encoding/base64" + "io" + "net/http" + "strings" + "testing" + + basechannels "github.com/sipeed/picoclaw/pkg/channels" + "github.com/sipeed/picoclaw/pkg/media" +) + +func TestStoreRemoteMedia_DetectsJPEGContentTypeFromBody(t *testing.T) { + t.Parallel() + + const jpegBase64 = "/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAP//////////////////////////////////////////////////////////////////////////////////////" + + "//////////////////////////////////////////////////////////////////////////////////////////////2wBDAf//////////////////////////////////////////////////////////////////////////////////////" + + "//////////////////////////////////////////////////////////////////////////////////////////////wAARCAABAAEDASIAAhEBAxEB/8QAFQABAQAAAAAAAAAAAAAAAAAAAAb/xAAVEQEBAAAAAAAAAAAAAAAAAAAABf/aAAwDAQACEAMQAAAB6A//xAAVEAEBAAAAAAAAAAAAAAAAAAAAEf/aAAgBAQABBQJf/8QAFBEBAAAAAAAAAAAAAAAAAAAAEP/aAAgBAwEBPwF//8QAFBEBAAAAAAAAAAAAAAAAAAAAEP/aAAgBAgEBPwF//8QAFBABAAAAAAAAAAAAAAAAAAAAEP/aAAgBAQAGPwJf/8QAFBABAAAAAAAAAAAAAAAAAAAAEP/aAAgBAQABPyFf/9k=" + + jpegData := decodeTestBase64(t, jpegBase64) + store := media.NewFileMediaStore() + ch := &WeComChannel{ + BaseChannel: basechannels.NewBaseChannel("wecom", nil, nil, nil), + mediaClient: &http.Client{ + Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/octet-stream"}}, + Body: io.NopCloser(bytes.NewReader(jpegData)), + }, nil + }), + }, + } + ch.SetMediaStore(store) + + ref, err := ch.storeRemoteMedia(context.Background(), "test-scope", "msg-1", "https://wecom.example/media", "", "") + if err != nil { + t.Fatalf("storeRemoteMedia returned error: %v", err) + } + t.Cleanup(func() { + _ = store.ReleaseAll("test-scope") + }) + + _, meta, err := store.ResolveWithMeta(ref) + if err != nil { + t.Fatalf("resolve media ref: %v", err) + } + if meta.ContentType != "image/jpeg" { + t.Fatalf("expected image/jpeg content type, got %q", meta.ContentType) + } + if !strings.HasSuffix(meta.Filename, ".jpg") && !strings.HasSuffix(meta.Filename, ".jpeg") { + t.Fatalf("expected jpeg filename, got %q", meta.Filename) + } +} + +func TestDetectWeComMediaMetadata_UsesFallbackExtensionWhenBodyUnknown(t *testing.T) { + t.Parallel() + + filename, contentType := detectWeComMediaMetadata([]byte("not a real image"), "msg-2.pdf", "", "", "") + if filename != "msg-2.pdf" { + t.Fatalf("expected fallback filename to be preserved, got %q", filename) + } + if contentType != "application/pdf" { + t.Fatalf("expected application/pdf from fallback extension, got %q", contentType) + } +} + +func TestStoreRemoteMedia_PreservesSuffixFromURL(t *testing.T) { + t.Parallel() + + docxLikeData := []byte("PK\x03\x04fake office payload") + store := media.NewFileMediaStore() + ch := &WeComChannel{ + BaseChannel: basechannels.NewBaseChannel("wecom", nil, nil, nil), + mediaClient: &http.Client{ + Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/octet-stream"}}, + Body: io.NopCloser(bytes.NewReader(docxLikeData)), + }, nil + }), + }, + } + ch.SetMediaStore(store) + + ref, err := ch.storeRemoteMedia( + context.Background(), + "test-scope", + "msg-docx", + "https://wecom.example/media/report.docx?signature=1", + "", + ".bin", + ) + if err != nil { + t.Fatalf("storeRemoteMedia returned error: %v", err) + } + t.Cleanup(func() { + _ = store.ReleaseAll("test-scope") + }) + + localPath, meta, err := store.ResolveWithMeta(ref) + if err != nil { + t.Fatalf("resolve media ref: %v", err) + } + if !strings.HasSuffix(meta.Filename, ".docx") { + t.Fatalf("expected docx filename, got %q", meta.Filename) + } + if !strings.HasSuffix(strings.ToLower(localPath), ".docx") { + t.Fatalf("expected docx temp path, got %q", localPath) + } +} + +func TestStoreRemoteMedia_PreservesSuffixFromContentDisposition(t *testing.T) { + t.Parallel() + + pptxLikeData := []byte("PK\x03\x04fake office payload") + store := media.NewFileMediaStore() + ch := &WeComChannel{ + BaseChannel: basechannels.NewBaseChannel("wecom", nil, nil, nil), + mediaClient: &http.Client{ + Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"application/octet-stream"}, + "Content-Disposition": []string{`attachment; filename="slides.pptx"`}, + }, + Body: io.NopCloser(bytes.NewReader(pptxLikeData)), + }, nil + }), + }, + } + ch.SetMediaStore(store) + + ref, err := ch.storeRemoteMedia( + context.Background(), + "test-scope", + "msg-pptx", + "https://wecom.example/media/download", + "", + ".bin", + ) + if err != nil { + t.Fatalf("storeRemoteMedia returned error: %v", err) + } + t.Cleanup(func() { + _ = store.ReleaseAll("test-scope") + }) + + localPath, meta, err := store.ResolveWithMeta(ref) + if err != nil { + t.Fatalf("resolve media ref: %v", err) + } + if !strings.HasSuffix(meta.Filename, ".pptx") { + t.Fatalf("expected pptx filename, got %q", meta.Filename) + } + if !strings.HasSuffix(strings.ToLower(localPath), ".pptx") { + t.Fatalf("expected pptx temp path, got %q", localPath) + } +} + +func decodeTestBase64(t *testing.T, value string) []byte { + t.Helper() + + data, err := io.ReadAll(base64.NewDecoder(base64.StdEncoding, strings.NewReader(value))) + if err != nil { + t.Fatalf("decode base64 fixture: %v", err) + } + return data +} + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { + return f(req) +} diff --git a/pkg/channels/wecom/protocol.go b/pkg/channels/wecom/protocol.go new file mode 100644 index 000000000..f42ce3bf4 --- /dev/null +++ b/pkg/channels/wecom/protocol.go @@ -0,0 +1,173 @@ +package wecom + +import "encoding/json" + +const ( + wecomDefaultWebSocketURL = "wss://openws.work.weixin.qq.com" + wecomCmdSubscribe = "aibot_subscribe" + wecomCmdPing = "ping" + wecomCmdMsgCallback = "aibot_msg_callback" + wecomCmdEventCallback = "aibot_event_callback" + wecomCmdRespondMsg = "aibot_respond_msg" + wecomCmdSendMsg = "aibot_send_msg" + wecomCmdUploadMediaInit = "aibot_upload_media_init" + wecomCmdUploadMediaChunk = "aibot_upload_media_chunk" + wecomCmdUploadMediaEnd = "aibot_upload_media_finish" +) + +type wecomEnvelope struct { + Cmd string `json:"cmd,omitempty"` + Headers wecomHeaders `json:"headers"` + Body json.RawMessage `json:"body,omitempty"` + ErrCode int `json:"errcode,omitempty"` + ErrMsg string `json:"errmsg,omitempty"` +} + +type wecomHeaders struct { + ReqID string `json:"req_id,omitempty"` +} + +type wecomCommand struct { + Cmd string `json:"cmd"` + Headers wecomHeaders `json:"headers"` + Body any `json:"body,omitempty"` +} + +type wecomSendMsgBody struct { + ChatID string `json:"chatid"` + ChatType uint32 `json:"chat_type,omitempty"` + MsgType string `json:"msgtype"` + Markdown *wecomMarkdownContent `json:"markdown,omitempty"` + File *wecomMediaRefContent `json:"file,omitempty"` + Image *wecomMediaRefContent `json:"image,omitempty"` + Voice *wecomMediaRefContent `json:"voice,omitempty"` + Video *wecomVideoContent `json:"video,omitempty"` + TemplateCard map[string]any `json:"template_card,omitempty"` +} + +type wecomRespondMsgBody struct { + MsgType string `json:"msgtype"` + Stream *wecomStreamContent `json:"stream,omitempty"` + Markdown *wecomMarkdownContent `json:"markdown,omitempty"` + File *wecomMediaRefContent `json:"file,omitempty"` + Image *wecomMediaRefContent `json:"image,omitempty"` + Voice *wecomMediaRefContent `json:"voice,omitempty"` + Video *wecomVideoContent `json:"video,omitempty"` + TemplateCard map[string]any `json:"template_card,omitempty"` +} + +type wecomStreamContent struct { + ID string `json:"id"` + Finish bool `json:"finish"` + Content string `json:"content,omitempty"` +} + +type wecomMarkdownContent struct { + Content string `json:"content"` +} + +type wecomMediaRefContent struct { + MediaID string `json:"media_id"` +} + +type wecomVideoContent struct { + MediaID string `json:"media_id"` + Title string `json:"title,omitempty"` + Description string `json:"description,omitempty"` +} + +type wecomUploadMediaInitBody struct { + Type string `json:"type"` + Filename string `json:"filename"` + TotalSize int64 `json:"total_size"` + TotalChunks int `json:"total_chunks"` + MD5 string `json:"md5,omitempty"` +} + +type wecomUploadMediaInitResponse struct { + UploadID string `json:"upload_id"` +} + +type wecomUploadMediaChunkBody struct { + UploadID string `json:"upload_id"` + ChunkIndex int `json:"chunk_index"` + Base64Data string `json:"base64_data"` +} + +type wecomUploadMediaFinishBody struct { + UploadID string `json:"upload_id"` +} + +type wecomUploadMediaFinishResponse struct { + Type string `json:"type"` + MediaID string `json:"media_id"` + CreatedAt json.RawMessage `json:"created_at"` +} + +type wecomIncomingMessage struct { + MsgID string `json:"msgid"` + AIBotID string `json:"aibotid"` + ChatID string `json:"chatid,omitempty"` + ChatType string `json:"chattype,omitempty"` + From struct { + UserID string `json:"userid"` + } `json:"from"` + MsgType string `json:"msgtype"` + Text *struct { + Content string `json:"content"` + } `json:"text,omitempty"` + Image *struct { + URL string `json:"url"` + AESKey string `json:"aeskey,omitempty"` + } `json:"image,omitempty"` + File *struct { + URL string `json:"url"` + AESKey string `json:"aeskey,omitempty"` + } `json:"file,omitempty"` + Video *struct { + URL string `json:"url"` + AESKey string `json:"aeskey,omitempty"` + } `json:"video,omitempty"` + Voice *struct { + Content string `json:"content"` + } `json:"voice,omitempty"` + Mixed *struct { + MsgItem []struct { + MsgType string `json:"msgtype"` + Text *struct { + Content string `json:"content"` + } `json:"text,omitempty"` + Image *struct { + URL string `json:"url"` + AESKey string `json:"aeskey,omitempty"` + } `json:"image,omitempty"` + File *struct { + URL string `json:"url"` + AESKey string `json:"aeskey,omitempty"` + } `json:"file,omitempty"` + } `json:"msg_item"` + } `json:"mixed,omitempty"` + Quote *struct { + MsgType string `json:"msgtype"` + Text *struct { + Content string `json:"content"` + } `json:"text,omitempty"` + } `json:"quote,omitempty"` + Event *struct { + EventType string `json:"eventtype"` + } `json:"event,omitempty"` +} + +func incomingChatID(msg wecomIncomingMessage) string { + if msg.ChatID != "" { + return msg.ChatID + } + return msg.From.UserID +} + +func incomingChatTypeCode(kind string) uint32 { + if kind == "group" { + return 2 + } + return 1 +} diff --git a/pkg/channels/wecom/reqid_store.go b/pkg/channels/wecom/reqid_store.go new file mode 100644 index 000000000..59e64e63d --- /dev/null +++ b/pkg/channels/wecom/reqid_store.go @@ -0,0 +1,113 @@ +package wecom + +import ( + "encoding/json" + "errors" + "os" + "path/filepath" + "sync" + "time" +) + +type wecomRoute struct { + ReqID string `json:"req_id"` + ChatID string `json:"chat_id"` + ChatType uint32 `json:"chat_type"` + ExpiresAt time.Time `json:"expires_at"` +} + +type reqIDStore struct { + mu sync.Mutex + path string + routes map[string]wecomRoute +} + +func newReqIDStore(path string) *reqIDStore { + if path == "" { + path = defaultReqIDStorePath() + } + s := &reqIDStore{ + path: path, + routes: make(map[string]wecomRoute), + } + _ = s.load() + return s +} + +func defaultReqIDStorePath() string { + if home, err := os.UserHomeDir(); err == nil && home != "" { + return filepath.Join(home, ".picoclaw", "wecom", "reqid-store.json") + } + return filepath.Join(os.TempDir(), "picoclaw-wecom-reqid-store.json") +} + +func (s *reqIDStore) Put(chatID, reqID string, chatType uint32, ttl time.Duration) error { + if reqID == "" || chatID == "" { + return nil + } + s.mu.Lock() + defer s.mu.Unlock() + s.deleteExpiredLocked(time.Now()) + s.routes[chatID] = wecomRoute{ + ReqID: reqID, + ChatID: chatID, + ChatType: chatType, + ExpiresAt: time.Now().Add(ttl), + } + return s.saveLocked() +} + +func (s *reqIDStore) Get(chatID string) (wecomRoute, bool) { + s.mu.Lock() + defer s.mu.Unlock() + s.deleteExpiredLocked(time.Now()) + route, ok := s.routes[chatID] + return route, ok +} + +func (s *reqIDStore) Delete(chatID string) error { + s.mu.Lock() + defer s.mu.Unlock() + delete(s.routes, chatID) + return s.saveLocked() +} + +func (s *reqIDStore) load() error { + s.mu.Lock() + defer s.mu.Unlock() + + data, err := os.ReadFile(s.path) + if err != nil { + if errors.Is(err, os.ErrNotExist) { + return nil + } + return err + } + + var routes map[string]wecomRoute + if err := json.Unmarshal(data, &routes); err != nil { + return err + } + s.routes = routes + s.deleteExpiredLocked(time.Now()) + return nil +} + +func (s *reqIDStore) deleteExpiredLocked(now time.Time) { + for chatID, route := range s.routes { + if !route.ExpiresAt.IsZero() && now.After(route.ExpiresAt) { + delete(s.routes, chatID) + } + } +} + +func (s *reqIDStore) saveLocked() error { + if err := os.MkdirAll(filepath.Dir(s.path), 0o700); err != nil { + return err + } + data, err := json.MarshalIndent(s.routes, "", " ") + if err != nil { + return err + } + return os.WriteFile(s.path, data, 0o600) +} diff --git a/pkg/channels/wecom/reqid_store_test.go b/pkg/channels/wecom/reqid_store_test.go new file mode 100644 index 000000000..e68e82500 --- /dev/null +++ b/pkg/channels/wecom/reqid_store_test.go @@ -0,0 +1,24 @@ +package wecom + +import ( + "path/filepath" + "testing" + "time" +) + +func TestReqIDStorePersistsRoutes(t *testing.T) { + storePath := filepath.Join(t.TempDir(), "reqids.json") + store := newReqIDStore(storePath) + if err := store.Put("chat-1", "req-1", 2, time.Hour); err != nil { + t.Fatalf("Put() error = %v", err) + } + + reloaded := newReqIDStore(storePath) + route, ok := reloaded.Get("chat-1") + if !ok { + t.Fatal("expected persisted route to be loaded") + } + if route.ChatID != "chat-1" || route.ReqID != "req-1" || route.ChatType != 2 { + t.Fatalf("loaded route = %+v", route) + } +} diff --git a/pkg/channels/wecom/wecom.go b/pkg/channels/wecom/wecom.go new file mode 100644 index 000000000..26e971921 --- /dev/null +++ b/pkg/channels/wecom/wecom.go @@ -0,0 +1,970 @@ +package wecom + +import ( + "context" + "crypto/rand" + "encoding/json" + "fmt" + "math/big" + "net/http" + "strings" + "sync" + "time" + + "github.com/gorilla/websocket" + + "github.com/sipeed/picoclaw/pkg/bus" + "github.com/sipeed/picoclaw/pkg/channels" + "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/identity" + "github.com/sipeed/picoclaw/pkg/logger" +) + +const ( + wecomConnectTimeout = 15 * time.Second + wecomCommandTimeout = 10 * time.Second + wecomUploadTimeout = 30 * time.Second + wecomHeartbeatInterval = 30 * time.Second + wecomStreamMaxDuration = 5*time.Minute + 30*time.Second + wecomStreamMinInterval = 500 * time.Millisecond + wecomRouteTTL = 30 * time.Minute + wecomMediaTimeout = 30 * time.Second + wecomRecentMessageMax = 1000 +) + +type WeComChannel struct { + *channels.BaseChannel + config config.WeComConfig + + ctx context.Context + cancel context.CancelFunc + + conn *websocket.Conn + connMu sync.Mutex + + pendingMu sync.Mutex + pending map[string]chan wecomEnvelope + + turnsMu sync.Mutex + turns map[string][]wecomTurn + + recent *recentMessageSet + routes *reqIDStore + mediaClient *http.Client + commandSend func(wecomCommand, time.Duration) (wecomEnvelope, error) +} + +type wecomTurn struct { + ReqID string + ChatID string + ChatType uint32 + StreamID string + CreatedAt time.Time +} + +type wecomStreamer struct { + channel *WeComChannel + chatID string + turn wecomTurn + + mu sync.Mutex + closed bool + lastSentAt time.Time + content string +} + +type recentMessageSet struct { + mu sync.Mutex + seen map[string]struct{} + ring []string + idx int +} + +func newRecentMessageSet(capacity int) *recentMessageSet { + if capacity <= 0 { + capacity = wecomRecentMessageMax + } + return &recentMessageSet{ + seen: make(map[string]struct{}, capacity), + ring: make([]string, capacity), + } +} + +func (s *recentMessageSet) Mark(id string) bool { + if id == "" { + return true + } + s.mu.Lock() + defer s.mu.Unlock() + if _, ok := s.seen[id]; ok { + return false + } + if old := s.ring[s.idx]; old != "" { + delete(s.seen, old) + } + s.ring[s.idx] = id + s.idx = (s.idx + 1) % len(s.ring) + s.seen[id] = struct{}{} + return true +} + +func NewChannel(cfg config.WeComConfig, messageBus *bus.MessageBus) (*WeComChannel, error) { + if cfg.BotID == "" || cfg.Secret() == "" { + return nil, fmt.Errorf("wecom bot_id and secret are required") + } + if cfg.WebSocketURL == "" { + cfg.WebSocketURL = wecomDefaultWebSocketURL + } + + base := channels.NewBaseChannel( + "wecom", + cfg, + messageBus, + cfg.AllowFrom, + channels.WithReasoningChannelID(cfg.ReasoningChannelID), + ) + + ch := &WeComChannel{ + BaseChannel: base, + config: cfg, + pending: make(map[string]chan wecomEnvelope), + turns: make(map[string][]wecomTurn), + recent: newRecentMessageSet(wecomRecentMessageMax), + routes: newReqIDStore(""), + mediaClient: &http.Client{Timeout: wecomMediaTimeout}, + } + ch.SetOwner(ch) + return ch, nil +} + +func (c *WeComChannel) Name() string { return "wecom" } + +func (c *WeComChannel) Start(ctx context.Context) error { + logger.InfoC("wecom", "Starting WeCom channel...") + c.ctx, c.cancel = context.WithCancel(ctx) + c.SetRunning(true) + go c.connectLoop() + return nil +} + +func (c *WeComChannel) Stop(_ context.Context) error { + logger.InfoC("wecom", "Stopping WeCom channel...") + if c.cancel != nil { + c.cancel() + } + c.connMu.Lock() + if c.conn != nil { + _ = c.conn.Close() + c.conn = nil + } + c.connMu.Unlock() + c.clearTurns() + c.SetRunning(false) + return nil +} + +func (c *WeComChannel) BeginStream(_ context.Context, chatID string) (channels.Streamer, error) { + if !c.IsRunning() { + return nil, channels.ErrNotRunning + } + + turn, ok := c.getTurn(chatID) + if !ok { + return nil, fmt.Errorf("wecom streaming unavailable: no active turn") + } + if time.Since(turn.CreatedAt) > wecomStreamMaxDuration { + c.consumeTurn(chatID, turn) + return nil, fmt.Errorf("wecom streaming unavailable: turn expired") + } + + return &wecomStreamer{ + channel: c, + chatID: chatID, + turn: turn, + }, nil +} + +func (c *WeComChannel) Send(ctx context.Context, msg bus.OutboundMessage) error { + if !c.IsRunning() { + return channels.ErrNotRunning + } + content := strings.TrimSpace(msg.Content) + if content == "" { + return nil + } + + if turn, ok := c.getTurn(msg.ChatID); ok { + if time.Since(turn.CreatedAt) <= wecomStreamMaxDuration { + if err := c.sendStreamReply(turn, content); err == nil { + c.consumeTurn(msg.ChatID, turn) + return nil + } + } + c.consumeTurn(msg.ChatID, turn) + } + + if route, ok := c.routes.Get(msg.ChatID); ok { + if err := c.sendActivePush(route.ChatID, route.ChatType, content); err != nil { + return err + } + return nil + } + + if err := c.sendActivePush(msg.ChatID, 0, content); err != nil { + return err + } + return nil +} + +func (c *WeComChannel) SendMedia(ctx context.Context, msg bus.OutboundMediaMessage) error { + if !c.IsRunning() { + return channels.ErrNotRunning + } + + route, chatType, hasTurn := c.resolveMediaRoute(msg.ChatID) + chatID := route.ChatID + if chatID == "" { + chatID = msg.ChatID + } + + for _, part := range msg.Parts { + if strings.TrimSpace(part.Ref) == "" { + if caption := strings.TrimSpace(part.Caption); caption != "" { + if err := c.sendActivePush(chatID, chatType, caption); err != nil { + return err + } + } + continue + } + + localPath, filename, contentType, cleanup, err := c.resolveOutboundPart(ctx, part) + if err != nil { + return fmt.Errorf("wecom resolve media %q: %v: %w", part.Ref, err, channels.ErrSendFailed) + } + + func() { + if cleanup != nil { + defer cleanup() + } + + uploaded, uploadErr := c.uploadOutboundMedia(ctx, localPath, filename, contentType, part) + if uploadErr != nil { + logger.WarnCF("wecom", "Falling back to placeholder after media upload failure", map[string]any{ + "chat_id": chatID, + "ref": part.Ref, + "filename": filename, + "content_type": contentType, + "error": uploadErr.Error(), + }) + if hasTurn { + if finishErr := c.sendStreamChunk(route, true, ""); finishErr != nil { + err = finishErr + return + } + c.deleteTurn(msg.ChatID) + hasTurn = false + } + err = c.sendActivePush(chatID, chatType, fallbackWeComMediaText(part, "", filename)) + return + } + + if hasTurn { + err = c.sendTurnMedia(route, uploaded) + c.deleteTurn(msg.ChatID) + hasTurn = false + } else { + err = c.sendActiveMedia(chatID, chatType, uploaded) + } + if err != nil { + return + } + if caption := strings.TrimSpace(part.Caption); caption != "" { + err = c.sendActivePush(chatID, chatType, caption) + } + }() + if err != nil { + return err + } + } + + return nil +} + +func (c *WeComChannel) connectLoop() { + backoff := time.Second + for { + select { + case <-c.ctx.Done(): + return + default: + } + + if err := c.runConnection(); err != nil { + logger.WarnCF("wecom", "WeCom connection lost", map[string]any{ + "error": err.Error(), + "backoff": backoff.String(), + }) + select { + case <-time.After(backoff): + case <-c.ctx.Done(): + return + } + if backoff < time.Minute { + backoff *= 2 + if backoff > time.Minute { + backoff = time.Minute + } + } + continue + } + return + } +} + +func (c *WeComChannel) runConnection() error { + dialCtx, cancel := context.WithTimeout(c.ctx, wecomConnectTimeout) + defer cancel() + + conn, resp, err := websocket.DefaultDialer.DialContext(dialCtx, c.config.WebSocketURL, nil) + if resp != nil { + _ = resp.Body.Close() + } + if err != nil { + return fmt.Errorf("%w: %v", channels.ErrTemporary, err) + } + + c.connMu.Lock() + c.conn = conn + c.connMu.Unlock() + defer func() { + c.connMu.Lock() + if c.conn == conn { + c.conn = nil + } + c.connMu.Unlock() + _ = conn.Close() + c.clearTurns() + }() + + readErrCh := make(chan error, 1) + go func() { + readErrCh <- c.readLoop(conn) + }() + + if writeErr := c.writeAndWait(conn, wecomCommand{ + Cmd: wecomCmdSubscribe, + Headers: wecomHeaders{ReqID: randomID(10)}, + Body: map[string]string{ + "bot_id": c.config.BotID, + "secret": c.config.Secret(), + }, + }, wecomCommandTimeout); writeErr != nil { + return writeErr + } + + heartbeatDone := make(chan struct{}) + go func() { + defer close(heartbeatDone) + c.heartbeatLoop(conn) + }() + + err = <-readErrCh + _ = conn.Close() + <-heartbeatDone + return err +} + +func (c *WeComChannel) heartbeatLoop(conn *websocket.Conn) { + ticker := time.NewTicker(wecomHeartbeatInterval) + defer ticker.Stop() + for { + select { + case <-ticker.C: + if err := c.writeAndWait(conn, wecomCommand{ + Cmd: wecomCmdPing, + Headers: wecomHeaders{ReqID: randomID(10)}, + }, wecomCommandTimeout); err != nil { + logger.WarnCF("wecom", "Heartbeat failed", map[string]any{"error": err.Error()}) + _ = conn.Close() + return + } + case <-c.ctx.Done(): + return + } + } +} + +func (c *WeComChannel) readLoop(conn *websocket.Conn) error { + for { + _, raw, err := conn.ReadMessage() + if err != nil { + select { + case <-c.ctx.Done(): + return nil + default: + return fmt.Errorf("%w: %v", channels.ErrTemporary, err) + } + } + + var env wecomEnvelope + if err := json.Unmarshal(raw, &env); err != nil { + logger.WarnCF("wecom", "Failed to parse WebSocket message", map[string]any{"error": err.Error()}) + continue + } + + if env.Cmd == "" && env.Headers.ReqID != "" { + c.pendingMu.Lock() + ch, ok := c.pending[env.Headers.ReqID] + if ok { + delete(c.pending, env.Headers.ReqID) + } + c.pendingMu.Unlock() + if ok { + ch <- env + } + continue + } + + go c.handleEnvelope(env) + } +} + +func (c *WeComChannel) handleEnvelope(env wecomEnvelope) { + switch env.Cmd { + case wecomCmdMsgCallback: + c.handleMessageCallback(env) + case wecomCmdEventCallback: + c.handleEventCallback(env) + default: + logger.DebugCF("wecom", "Ignoring unsupported WeCom command", map[string]any{"cmd": env.Cmd}) + } +} + +func (c *WeComChannel) handleEventCallback(env wecomEnvelope) { + var msg wecomIncomingMessage + if err := json.Unmarshal(env.Body, &msg); err != nil { + logger.WarnCF("wecom", "Failed to parse WeCom event callback", map[string]any{"error": err.Error()}) + } +} + +func (c *WeComChannel) handleMessageCallback(env wecomEnvelope) { + var msg wecomIncomingMessage + if err := json.Unmarshal(env.Body, &msg); err != nil { + logger.WarnCF("wecom", "Failed to parse WeCom message callback", map[string]any{"error": err.Error()}) + return + } + if !c.recent.Mark(msg.MsgID) { + return + } + + reqID := env.Headers.ReqID + if reqID == "" { + logger.WarnC("wecom", "WeCom message callback missing req_id") + return + } + if msg.Event != nil && msg.Event.EventType != "" { + return + } + + if err := c.dispatchIncoming(reqID, msg); err != nil { + logger.WarnCF("wecom", "Failed to dispatch WeCom message", map[string]any{ + "req_id": reqID, + "error": err.Error(), + }) + _ = c.respondImmediate(reqID, "The WeCom message could not be processed.") + } +} + +func (c *WeComChannel) dispatchIncoming(reqID string, msg wecomIncomingMessage) error { + senderID := msg.From.UserID + if senderID == "" { + senderID = "unknown" + } + actualChatID := incomingChatID(msg) + chatType := incomingChatTypeCode(msg.ChatType) + peerKind := "direct" + if msg.ChatType == "group" { + peerKind = "group" + } + + sender := bus.SenderInfo{ + Platform: "wecom", + PlatformID: senderID, + CanonicalID: identity.BuildCanonicalID("wecom", senderID), + DisplayName: senderID, + } + + var ( + content string + quoteText string + mediaRefs []string + err error + ) + scope := channels.BuildMediaScope("wecom", actualChatID, msg.MsgID) + switch msg.MsgType { + case "text": + if msg.Text != nil { + content = strings.TrimSpace(msg.Text.Content) + } + case "voice": + if msg.Voice != nil { + content = strings.TrimSpace(msg.Voice.Content) + } + case "image": + content = "[image]" + mediaRefs, err = c.collectSingleMedia(c.ctx, scope, msg.MsgID, &mediaPayload{ + url: msg.Image.URL, + aesKey: msg.Image.AESKey, + }, "image", ".jpg") + case "file": + content = "[file]" + mediaRefs, err = c.collectSingleMedia(c.ctx, scope, msg.MsgID, &mediaPayload{ + url: msg.File.URL, + aesKey: msg.File.AESKey, + }, "file", ".bin") + case "video": + content = "[video]" + mediaRefs, err = c.collectSingleMedia(c.ctx, scope, msg.MsgID, &mediaPayload{ + url: msg.Video.URL, + aesKey: msg.Video.AESKey, + }, "video", ".mp4") + case "mixed": + content, mediaRefs, err = c.collectMixedMedia(c.ctx, scope, msg) + default: + return c.respondImmediate(reqID, "Unsupported WeCom message type: "+msg.MsgType) + } + if err != nil { + return err + } + if msg.Quote != nil && msg.Quote.Text != nil { + quoteText = strings.TrimSpace(msg.Quote.Text.Content) + if content == "" { + content = quoteText + } + } + if content == "" && len(mediaRefs) == 0 { + return c.respondImmediate(reqID, "The WeCom message did not contain usable content.") + } + + turn := wecomTurn{ + ReqID: reqID, + ChatID: actualChatID, + ChatType: chatType, + StreamID: randomID(10), + CreatedAt: time.Now(), + } + c.queueTurn(actualChatID, turn) + if err := c.routes.Put(actualChatID, reqID, chatType, wecomRouteTTL); err != nil { + logger.WarnCF("wecom", "Failed to persist req_id route", map[string]any{ + "chat_id": actualChatID, + "req_id": reqID, + "error": err.Error(), + }) + } + + opening := "" + if c.config.SendThinkingMessage { + opening = "Processing..." + } + if err := c.sendStreamChunk(turn, false, opening); err != nil { + return err + } + + peer := bus.Peer{Kind: peerKind, ID: actualChatID} + metadata := map[string]string{ + "channel": "wecom", + "req_id": reqID, + "chat_id": actualChatID, + "chat_type": msg.ChatType, + "msg_id": msg.MsgID, + "msg_type": msg.MsgType, + } + if quoteText != "" { + metadata["quote_text"] = quoteText + } + + c.HandleMessage(c.ctx, peer, msg.MsgID, senderID, actualChatID, content, mediaRefs, metadata, sender) + return nil +} + +func (c *WeComChannel) collectSingleMedia( + ctx context.Context, + scope, msgID string, + payload interface { + GetURL() string + GetAESKey() string + }, + label, fallbackExt string, +) ([]string, error) { + if payload == nil || payload.GetURL() == "" { + return nil, fmt.Errorf("%s payload is empty", label) + } + ref, err := c.storeRemoteMedia(ctx, scope, msgID, payload.GetURL(), payload.GetAESKey(), fallbackExt) + if err != nil { + return nil, err + } + return []string{ref}, nil +} + +type mediaPayload struct { + url string + aesKey string +} + +func (p *mediaPayload) GetURL() string { return p.url } +func (p *mediaPayload) GetAESKey() string { return p.aesKey } + +func (c *WeComChannel) collectMixedMedia( + ctx context.Context, + scope string, + msg wecomIncomingMessage, +) (string, []string, error) { + if msg.Mixed == nil { + return "", nil, fmt.Errorf("mixed message is empty") + } + + var textParts []string + var refs []string + for idx, item := range msg.Mixed.MsgItem { + switch item.MsgType { + case "text": + if item.Text != nil && strings.TrimSpace(item.Text.Content) != "" { + textParts = append(textParts, strings.TrimSpace(item.Text.Content)) + } + case "image": + if item.Image != nil && item.Image.URL != "" { + ref, err := c.storeRemoteMedia( + ctx, + scope, + fmt.Sprintf("%s-%d", msg.MsgID, idx), + item.Image.URL, + item.Image.AESKey, + ".jpg", + ) + if err != nil { + return "", nil, err + } + refs = append(refs, ref) + } + case "file": + if item.File != nil && item.File.URL != "" { + ref, err := c.storeRemoteMedia( + ctx, + scope, + fmt.Sprintf("%s-%d", msg.MsgID, idx), + item.File.URL, + item.File.AESKey, + ".bin", + ) + if err != nil { + return "", nil, err + } + refs = append(refs, ref) + } + } + } + + content := strings.Join(textParts, "\n") + if content == "" && len(refs) > 0 { + content = "[media]" + } + return content, refs, nil +} + +func (c *WeComChannel) respondImmediate(reqID, content string) error { + turn := wecomTurn{ + ReqID: reqID, + StreamID: randomID(10), + CreatedAt: time.Now(), + } + return c.sendStreamChunk(turn, true, content) +} + +func (c *WeComChannel) sendStreamReply(turn wecomTurn, content string) error { + return c.sendStreamChunk(turn, true, content) +} + +func (c *WeComChannel) sendStreamChunk(turn wecomTurn, finish bool, content string) error { + return c.sendCommand(wecomCommand{ + Cmd: wecomCmdRespondMsg, + Headers: wecomHeaders{ReqID: turn.ReqID}, + Body: wecomRespondMsgBody{ + MsgType: "stream", + Stream: &wecomStreamContent{ + ID: turn.StreamID, + Finish: finish, + Content: content, + }, + }, + }, wecomCommandTimeout) +} + +func (c *WeComChannel) sendTurnMedia(turn wecomTurn, uploaded *wecomOutboundMedia) error { + if uploaded == nil { + return fmt.Errorf("wecom outbound media is nil: %w", channels.ErrSendFailed) + } + if err := c.sendCommand(wecomCommand{ + Cmd: wecomCmdRespondMsg, + Headers: wecomHeaders{ReqID: turn.ReqID}, + Body: uploaded.respondBody(), + }, wecomCommandTimeout); err != nil { + return err + } + return c.sendStreamChunk(turn, true, "") +} + +func (c *WeComChannel) sendActivePush(chatID string, chatType uint32, content string) error { + if strings.TrimSpace(chatID) == "" { + return fmt.Errorf("empty chat ID: %w", channels.ErrSendFailed) + } + return c.sendCommand(wecomCommand{ + Cmd: wecomCmdSendMsg, + Headers: wecomHeaders{ReqID: randomID(10)}, + Body: wecomSendMsgBody{ + ChatID: chatID, + ChatType: chatType, + MsgType: "markdown", + Markdown: &wecomMarkdownContent{Content: content}, + }, + }, wecomCommandTimeout) +} + +func (c *WeComChannel) sendActiveMedia(chatID string, chatType uint32, uploaded *wecomOutboundMedia) error { + if strings.TrimSpace(chatID) == "" { + return fmt.Errorf("empty chat ID: %w", channels.ErrSendFailed) + } + if uploaded == nil { + return fmt.Errorf("wecom outbound media is nil: %w", channels.ErrSendFailed) + } + return c.sendCommand(wecomCommand{ + Cmd: wecomCmdSendMsg, + Headers: wecomHeaders{ReqID: randomID(10)}, + Body: uploaded.sendBody(chatID, chatType), + }, wecomCommandTimeout) +} + +func (c *WeComChannel) sendCommand(cmd wecomCommand, timeout time.Duration) error { + _, err := c.sendCommandAck(cmd, timeout) + return err +} + +func (c *WeComChannel) sendCommandAck(cmd wecomCommand, timeout time.Duration) (wecomEnvelope, error) { + if c.commandSend != nil { + return c.commandSend(cmd, timeout) + } + return c.writeCurrentAck(cmd, timeout) +} + +func (c *WeComChannel) writeCurrentAck(cmd wecomCommand, timeout time.Duration) (wecomEnvelope, error) { + c.connMu.Lock() + conn := c.conn + c.connMu.Unlock() + if conn == nil { + return wecomEnvelope{}, fmt.Errorf("wecom websocket not connected: %w", channels.ErrTemporary) + } + return c.writeAndWaitAck(conn, cmd, timeout) +} + +func (c *WeComChannel) writeAndWait(conn *websocket.Conn, cmd wecomCommand, timeout time.Duration) error { + _, err := c.writeAndWaitAck(conn, cmd, timeout) + return err +} + +func (c *WeComChannel) writeAndWaitAck( + conn *websocket.Conn, + cmd wecomCommand, + timeout time.Duration, +) (wecomEnvelope, error) { + if cmd.Headers.ReqID == "" { + cmd.Headers.ReqID = randomID(10) + } + waitCh := make(chan wecomEnvelope, 1) + c.pendingMu.Lock() + c.pending[cmd.Headers.ReqID] = waitCh + c.pendingMu.Unlock() + defer func() { + c.pendingMu.Lock() + delete(c.pending, cmd.Headers.ReqID) + c.pendingMu.Unlock() + }() + + data, err := json.Marshal(cmd) + if err != nil { + return wecomEnvelope{}, fmt.Errorf("%w: %v", channels.ErrSendFailed, err) + } + c.connMu.Lock() + err = conn.WriteMessage(websocket.TextMessage, data) + c.connMu.Unlock() + if err != nil { + return wecomEnvelope{}, fmt.Errorf("%w: %v", channels.ErrTemporary, err) + } + + timer := time.NewTimer(timeout) + defer timer.Stop() + select { + case env := <-waitCh: + if env.ErrCode != 0 { + return wecomEnvelope{}, fmt.Errorf( + "%w: wecom errcode=%d errmsg=%s", + channels.ErrTemporary, + env.ErrCode, + env.ErrMsg, + ) + } + return env, nil + case <-timer.C: + return wecomEnvelope{}, fmt.Errorf("%w: timeout waiting for WeCom ack", channels.ErrTemporary) + case <-c.ctx.Done(): + return wecomEnvelope{}, c.ctx.Err() + } +} + +func (c *WeComChannel) getTurn(chatID string) (wecomTurn, bool) { + c.turnsMu.Lock() + defer c.turnsMu.Unlock() + queue := c.turns[chatID] + if len(queue) == 0 { + return wecomTurn{}, false + } + return queue[0], true +} + +func (c *WeComChannel) deleteTurn(chatID string) { + c.turnsMu.Lock() + defer c.turnsMu.Unlock() + queue := c.turns[chatID] + if len(queue) <= 1 { + delete(c.turns, chatID) + return + } + c.turns[chatID] = queue[1:] +} + +func (c *WeComChannel) queueTurn(chatID string, turn wecomTurn) { + c.turnsMu.Lock() + defer c.turnsMu.Unlock() + c.turns[chatID] = append(c.turns[chatID], turn) +} + +func (c *WeComChannel) consumeTurn(chatID string, turn wecomTurn) bool { + c.turnsMu.Lock() + defer c.turnsMu.Unlock() + + queue := c.turns[chatID] + if len(queue) == 0 { + return false + } + current := queue[0] + if current.ReqID != turn.ReqID || current.StreamID != turn.StreamID { + return false + } + if len(queue) == 1 { + delete(c.turns, chatID) + return true + } + c.turns[chatID] = queue[1:] + return true +} + +func (c *WeComChannel) clearTurns() { + c.turnsMu.Lock() + c.turns = make(map[string][]wecomTurn) + c.turnsMu.Unlock() +} + +func randomID(n int) string { + const alphabet = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789" + if n <= 0 { + n = 10 + } + buf := make([]byte, n) + for i := range buf { + v, _ := rand.Int(rand.Reader, big.NewInt(int64(len(alphabet)))) + buf[i] = alphabet[v.Int64()] + } + return string(buf) +} + +func (s *wecomStreamer) Update(ctx context.Context, content string) error { + s.mu.Lock() + defer s.mu.Unlock() + + if s.closed { + return nil + } + if err := s.validateActiveTurn(); err != nil { + return err + } + if err := ctx.Err(); err != nil { + return err + } + + if !s.lastSentAt.IsZero() { + wait := time.Until(s.lastSentAt.Add(wecomStreamMinInterval)) + if wait > 0 { + timer := time.NewTimer(wait) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + } + } + } + + if err := s.channel.sendStreamChunk(s.turn, false, content); err != nil { + return err + } + s.content = content + s.lastSentAt = time.Now() + return nil +} + +func (s *wecomStreamer) Finalize(ctx context.Context, content string) error { + s.mu.Lock() + defer s.mu.Unlock() + + if s.closed { + return nil + } + if err := s.validateActiveTurn(); err != nil { + return err + } + if err := ctx.Err(); err != nil { + return err + } + if err := s.channel.sendStreamChunk(s.turn, true, content); err != nil { + return err + } + + s.content = content + s.closed = true + s.channel.consumeTurn(s.chatID, s.turn) + return nil +} + +func (s *wecomStreamer) Cancel(_ context.Context) { + s.mu.Lock() + defer s.mu.Unlock() + + if s.closed { + return + } + if s.validateActiveTurn() == nil { + _ = s.channel.sendStreamChunk(s.turn, true, s.content) + s.channel.consumeTurn(s.chatID, s.turn) + } + s.closed = true +} + +func (s *wecomStreamer) validateActiveTurn() error { + if time.Since(s.turn.CreatedAt) > wecomStreamMaxDuration { + s.channel.consumeTurn(s.chatID, s.turn) + return fmt.Errorf("wecom streaming unavailable: turn expired") + } + current, ok := s.channel.getTurn(s.chatID) + if !ok || current.ReqID != s.turn.ReqID || current.StreamID != s.turn.StreamID { + return fmt.Errorf("wecom streaming unavailable: turn no longer active") + } + return nil +} diff --git a/pkg/channels/wecom/wecom_test.go b/pkg/channels/wecom/wecom_test.go new file mode 100644 index 000000000..c7a4adfc0 --- /dev/null +++ b/pkg/channels/wecom/wecom_test.go @@ -0,0 +1,660 @@ +package wecom + +import ( + "context" + "encoding/json" + "errors" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/sipeed/picoclaw/pkg/bus" + "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/media" +) + +func TestDispatchIncoming_UsesActualChatIDAndStoresReqIDRoute(t *testing.T) { + t.Parallel() + + messageBus := bus.NewMessageBus() + ch := newTestWeComChannel(t, messageBus) + + var commands []wecomCommand + ch.commandSend = func(cmd wecomCommand, _ time.Duration) (wecomEnvelope, error) { + commands = append(commands, cmd) + return wecomTestAck(nil), nil + } + + msg := wecomIncomingMessage{ + MsgID: "msg-1", + ChatID: "chat-1", + ChatType: "direct", + MsgType: "text", + Text: &struct { + Content string `json:"content"` + }{Content: "hello"}, + } + msg.From.UserID = "user-1" + + if err := ch.dispatchIncoming("req-1", msg); err != nil { + t.Fatalf("dispatchIncoming() error = %v", err) + } + + select { + case inbound := <-messageBus.InboundChan(): + if inbound.ChatID != "chat-1" { + t.Fatalf("inbound ChatID = %q, want chat-1", inbound.ChatID) + } + if inbound.MessageID != "msg-1" { + t.Fatalf("inbound MessageID = %q, want msg-1", inbound.MessageID) + } + if inbound.Peer.ID != "chat-1" { + t.Fatalf("inbound Peer.ID = %q, want chat-1", inbound.Peer.ID) + } + if inbound.Metadata["req_id"] != "req-1" { + t.Fatalf("inbound req_id = %q, want req-1", inbound.Metadata["req_id"]) + } + default: + t.Fatal("expected inbound message to be published") + } + + turn, ok := ch.getTurn("chat-1") + if !ok { + t.Fatal("expected queued turn for chat-1") + } + if turn.ReqID != "req-1" { + t.Fatalf("turn.ReqID = %q, want req-1", turn.ReqID) + } + + route, ok := ch.routes.Get("chat-1") + if !ok { + t.Fatal("expected persisted route for chat-1") + } + if route.ReqID != "req-1" || route.ChatType != 1 { + t.Fatalf("route = %+v", route) + } + + if len(commands) != 1 { + t.Fatalf("expected 1 opening command, got %d", len(commands)) + } + if commands[0].Cmd != wecomCmdRespondMsg { + t.Fatalf("opening command = %q, want %q", commands[0].Cmd, wecomCmdRespondMsg) + } + if commands[0].Headers.ReqID != "req-1" { + t.Fatalf("opening req_id = %q, want req-1", commands[0].Headers.ReqID) + } +} + +func TestNewChannel_DoesNotRegisterMessageSplitLimit(t *testing.T) { + t.Parallel() + + ch := newTestWeComChannel(t, bus.NewMessageBus()) + if got := ch.MaxMessageLength(); got != 0 { + t.Fatalf("MaxMessageLength() = %d, want 0", got) + } +} + +func TestBeginStream_UpdateAndFinalize(t *testing.T) { + t.Parallel() + + ch := newTestWeComChannel(t, bus.NewMessageBus()) + ch.SetRunning(true) + ch.queueTurn("chat-1", wecomTurn{ + ReqID: "req-1", + ChatID: "chat-1", + ChatType: 1, + StreamID: "stream-1", + CreatedAt: time.Now(), + }) + + var commands []wecomCommand + ch.commandSend = func(cmd wecomCommand, _ time.Duration) (wecomEnvelope, error) { + commands = append(commands, cmd) + return wecomTestAck(nil), nil + } + + streamer, err := ch.BeginStream(context.Background(), "chat-1") + if err != nil { + t.Fatalf("BeginStream() error = %v", err) + } + if err := streamer.Update(context.Background(), "draft"); err != nil { + t.Fatalf("Update() error = %v", err) + } + if err := streamer.Finalize(context.Background(), "final"); err != nil { + t.Fatalf("Finalize() error = %v", err) + } + + if len(commands) != 2 { + t.Fatalf("expected 2 commands, got %d", len(commands)) + } + for i, wantFinish := range []bool{false, true} { + if commands[i].Cmd != wecomCmdRespondMsg { + t.Fatalf("command[%d].Cmd = %q, want %q", i, commands[i].Cmd, wecomCmdRespondMsg) + } + body, ok := commands[i].Body.(wecomRespondMsgBody) + if !ok { + t.Fatalf("command[%d] body type = %T", i, commands[i].Body) + } + if body.Stream == nil { + t.Fatalf("command[%d] missing stream body", i) + } + if body.Stream.ID != "stream-1" { + t.Fatalf("command[%d] stream id = %q, want stream-1", i, body.Stream.ID) + } + if body.Stream.Finish != wantFinish { + t.Fatalf("command[%d] finish = %v, want %v", i, body.Stream.Finish, wantFinish) + } + } + if body := commands[0].Body.(wecomRespondMsgBody); body.Stream.Content != "draft" { + t.Fatalf("update content = %q, want draft", body.Stream.Content) + } + if body := commands[1].Body.(wecomRespondMsgBody); body.Stream.Content != "final" { + t.Fatalf("final content = %q, want final", body.Stream.Content) + } + if _, ok := ch.getTurn("chat-1"); ok { + t.Fatal("expected turn to be consumed after Finalize") + } +} + +func TestSend_StreamFailureFallsBackToActualChatID(t *testing.T) { + t.Parallel() + + ch := newTestWeComChannel(t, bus.NewMessageBus()) + ch.SetRunning(true) + ch.queueTurn("chat-1", wecomTurn{ + ReqID: "req-1", + ChatID: "chat-1", + ChatType: 1, + StreamID: "stream-1", + CreatedAt: time.Now(), + }) + ch.queueTurn("chat-1", wecomTurn{ + ReqID: "req-2", + ChatID: "chat-1", + ChatType: 1, + StreamID: "stream-2", + CreatedAt: time.Now(), + }) + if err := ch.routes.Put("chat-1", "req-2", 1, time.Hour); err != nil { + t.Fatalf("Put() error = %v", err) + } + + var commands []wecomCommand + ch.commandSend = func(cmd wecomCommand, _ time.Duration) (wecomEnvelope, error) { + commands = append(commands, cmd) + if len(commands) == 1 && cmd.Cmd == wecomCmdRespondMsg { + return wecomEnvelope{}, errors.New("stream send failed") + } + return wecomTestAck(nil), nil + } + + if err := ch.Send(context.Background(), bus.OutboundMessage{ + Channel: "wecom", + ChatID: "chat-1", + Content: "hello", + }); err != nil { + t.Fatalf("Send() error = %v", err) + } + + if len(commands) != 2 { + t.Fatalf("expected 2 commands, got %d", len(commands)) + } + if commands[0].Cmd != wecomCmdRespondMsg || commands[0].Headers.ReqID != "req-1" { + t.Fatalf("first command = %+v", commands[0]) + } + if commands[1].Cmd != wecomCmdSendMsg { + t.Fatalf("second command = %q, want %q", commands[1].Cmd, wecomCmdSendMsg) + } + body, ok := commands[1].Body.(wecomSendMsgBody) + if !ok { + t.Fatalf("unexpected send body type %T", commands[1].Body) + } + if body.ChatID != "chat-1" { + t.Fatalf("send chatid = %q, want chat-1", body.ChatID) + } + if body.ChatType != 1 { + t.Fatalf("send chat_type = %d, want 1", body.ChatType) + } + + nextTurn, ok := ch.getTurn("chat-1") + if !ok { + t.Fatal("expected second turn to remain queued") + } + if nextTurn.ReqID != "req-2" { + t.Fatalf("next queued req_id = %q, want req-2", nextTurn.ReqID) + } +} + +func TestSend_DoesNotSplitStreamReply(t *testing.T) { + t.Parallel() + + ch := newTestWeComChannel(t, bus.NewMessageBus()) + ch.SetRunning(true) + ch.queueTurn("chat-1", wecomTurn{ + ReqID: "req-1", + ChatID: "chat-1", + ChatType: 1, + StreamID: "stream-1", + CreatedAt: time.Now(), + }) + + var commands []wecomCommand + ch.commandSend = func(cmd wecomCommand, _ time.Duration) (wecomEnvelope, error) { + commands = append(commands, cmd) + return wecomTestAck(nil), nil + } + + content := strings.Repeat("\u4e2d", 30000) + if err := ch.Send(context.Background(), bus.OutboundMessage{ + Channel: "wecom", + ChatID: "chat-1", + Content: content, + }); err != nil { + t.Fatalf("Send() error = %v", err) + } + + if len(commands) != 1 { + t.Fatalf("expected 1 stream command, got %d", len(commands)) + } + body, ok := commands[0].Body.(wecomRespondMsgBody) + if !ok { + t.Fatalf("unexpected body type %T", commands[0].Body) + } + if body.Stream == nil || !body.Stream.Finish { + t.Fatalf("stream body = %+v", body.Stream) + } + if body.Stream.Content != content { + t.Fatalf("stream content length = %d, want %d", len(body.Stream.Content), len(content)) + } +} + +func TestSend_DoesNotSplitActivePush(t *testing.T) { + t.Parallel() + + ch := newTestWeComChannel(t, bus.NewMessageBus()) + ch.SetRunning(true) + + var commands []wecomCommand + ch.commandSend = func(cmd wecomCommand, _ time.Duration) (wecomEnvelope, error) { + commands = append(commands, cmd) + return wecomTestAck(nil), nil + } + + content := strings.Repeat("a", 30000) + if err := ch.Send(context.Background(), bus.OutboundMessage{ + Channel: "wecom", + ChatID: "chat-1", + Content: content, + }); err != nil { + t.Fatalf("Send() error = %v", err) + } + + if len(commands) != 1 { + t.Fatalf("expected 1 send command, got %d", len(commands)) + } + if commands[0].Cmd != wecomCmdSendMsg { + t.Fatalf("command = %q, want %q", commands[0].Cmd, wecomCmdSendMsg) + } + body, ok := commands[0].Body.(wecomSendMsgBody) + if !ok { + t.Fatalf("unexpected body type %T", commands[0].Body) + } + if body.Markdown == nil || body.Markdown.Content != content { + t.Fatalf("markdown content length = %d, want %d", len(body.Markdown.Content), len(content)) + } +} + +func TestSendMedia_SendsActiveImage(t *testing.T) { + t.Parallel() + + ch := newTestWeComChannel(t, bus.NewMessageBus()) + ch.SetRunning(true) + + store := media.NewFileMediaStore() + ch.SetMediaStore(store) + + imageData := wecomTestJPEGData(t) + imagePath := filepath.Join(t.TempDir(), "photo.jpg") + if err := os.WriteFile(imagePath, imageData, 0o600); err != nil { + t.Fatalf("WriteFile() error = %v", err) + } + ref, err := store.Store(imagePath, media.MediaMeta{ + Filename: "photo.jpg", + ContentType: "image/jpeg", + Source: "test", + CleanupPolicy: media.CleanupPolicyForgetOnly, + }, "scope-1") + if err != nil { + t.Fatalf("Store() error = %v", err) + } + + var commands []wecomCommand + ch.commandSend = func(cmd wecomCommand, _ time.Duration) (wecomEnvelope, error) { + commands = append(commands, cmd) + switch cmd.Cmd { + case wecomCmdUploadMediaInit: + return wecomTestAck(wecomUploadMediaInitResponse{UploadID: "upload-1"}), nil + case wecomCmdUploadMediaEnd: + return wecomTestAck(wecomUploadMediaFinishResponse{ + Type: "image", + MediaID: "media-1", + }), nil + default: + return wecomTestAck(nil), nil + } + } + + err = ch.SendMedia(context.Background(), bus.OutboundMediaMessage{ + Channel: "wecom", + ChatID: "chat-1", + Parts: []bus.MediaPart{{ + Ref: ref, + Type: "image", + Filename: "photo.jpg", + ContentType: "image/jpeg", + }}, + }) + if err != nil { + t.Fatalf("SendMedia() error = %v", err) + } + + if len(commands) != 4 { + t.Fatalf("expected 4 commands, got %d", len(commands)) + } + if commands[0].Cmd != wecomCmdUploadMediaInit { + t.Fatalf("first command = %q, want %q", commands[0].Cmd, wecomCmdUploadMediaInit) + } + initBody, ok := commands[0].Body.(wecomUploadMediaInitBody) + if !ok { + t.Fatalf("unexpected init body type %T", commands[0].Body) + } + if initBody.Type != "image" || initBody.Filename != "photo.jpg" || initBody.TotalChunks != 1 { + t.Fatalf("init body = %+v", initBody) + } + if commands[1].Cmd != wecomCmdUploadMediaChunk { + t.Fatalf("second command = %q, want %q", commands[1].Cmd, wecomCmdUploadMediaChunk) + } + chunkBody, ok := commands[1].Body.(wecomUploadMediaChunkBody) + if !ok { + t.Fatalf("unexpected chunk body type %T", commands[1].Body) + } + if chunkBody.UploadID != "upload-1" || chunkBody.ChunkIndex != 0 || chunkBody.Base64Data == "" { + t.Fatalf("chunk body = %+v", chunkBody) + } + if commands[2].Cmd != wecomCmdUploadMediaEnd { + t.Fatalf("third command = %q, want %q", commands[2].Cmd, wecomCmdUploadMediaEnd) + } + if commands[3].Cmd != wecomCmdSendMsg { + t.Fatalf("fourth command = %q, want %q", commands[3].Cmd, wecomCmdSendMsg) + } + + body, ok := commands[3].Body.(wecomSendMsgBody) + if !ok { + t.Fatalf("unexpected send body type %T", commands[3].Body) + } + if body.MsgType != "image" || body.Image == nil { + t.Fatalf("send body = %+v", body) + } + if body.ChatID != "chat-1" { + t.Fatalf("send chatid = %q, want chat-1", body.ChatID) + } + if body.Image.MediaID != "media-1" { + t.Fatalf("image media_id = %q, want media-1", body.Image.MediaID) + } +} + +func TestSendMedia_UsesTurnImageAndFinishesStream(t *testing.T) { + t.Parallel() + + ch := newTestWeComChannel(t, bus.NewMessageBus()) + ch.SetRunning(true) + + store := media.NewFileMediaStore() + ch.SetMediaStore(store) + + imageData := wecomTestJPEGData(t) + imagePath := filepath.Join(t.TempDir(), "reply.jpg") + if err := os.WriteFile(imagePath, imageData, 0o600); err != nil { + t.Fatalf("WriteFile() error = %v", err) + } + ref, err := store.Store(imagePath, media.MediaMeta{ + Filename: "reply.jpg", + ContentType: "image/jpeg", + Source: "test", + CleanupPolicy: media.CleanupPolicyForgetOnly, + }, "scope-2") + if err != nil { + t.Fatalf("Store() error = %v", err) + } + + ch.queueTurn("chat-1", wecomTurn{ + ReqID: "req-1", + ChatID: "chat-1", + ChatType: 1, + StreamID: "stream-1", + CreatedAt: time.Now(), + }) + putErr := ch.routes.Put("chat-1", "req-1", 1, time.Hour) + if putErr != nil { + t.Fatalf("Put() error = %v", putErr) + } + + var commands []wecomCommand + ch.commandSend = func(cmd wecomCommand, _ time.Duration) (wecomEnvelope, error) { + commands = append(commands, cmd) + switch cmd.Cmd { + case wecomCmdUploadMediaInit: + return wecomTestAck(wecomUploadMediaInitResponse{UploadID: "upload-2"}), nil + case wecomCmdUploadMediaEnd: + return wecomTestAck(wecomUploadMediaFinishResponse{ + Type: "image", + MediaID: "media-2", + }), nil + default: + return wecomTestAck(nil), nil + } + } + + err = ch.SendMedia(context.Background(), bus.OutboundMediaMessage{ + Channel: "wecom", + ChatID: "chat-1", + Parts: []bus.MediaPart{{ + Ref: ref, + Type: "image", + Filename: "reply.jpg", + ContentType: "image/jpeg", + }}, + }) + if err != nil { + t.Fatalf("SendMedia() error = %v", err) + } + + if len(commands) != 5 { + t.Fatalf("expected 5 commands, got %d", len(commands)) + } + if commands[0].Cmd != wecomCmdUploadMediaInit { + t.Fatalf("first command = %+v", commands[0]) + } + if commands[1].Cmd != wecomCmdUploadMediaChunk { + t.Fatalf("second command = %+v", commands[1]) + } + if commands[2].Cmd != wecomCmdUploadMediaEnd { + t.Fatalf("third command = %+v", commands[2]) + } + if commands[3].Cmd != wecomCmdRespondMsg || commands[3].Headers.ReqID != "req-1" { + t.Fatalf("fourth command = %+v", commands[3]) + } + if commands[4].Cmd != wecomCmdRespondMsg || commands[4].Headers.ReqID != "req-1" { + t.Fatalf("fifth command = %+v", commands[4]) + } + + imageBody, ok := commands[3].Body.(wecomRespondMsgBody) + if !ok { + t.Fatalf("unexpected image body type %T", commands[3].Body) + } + if imageBody.MsgType != "image" || imageBody.Image == nil { + t.Fatalf("image body = %+v", imageBody) + } + if imageBody.Image.MediaID != "media-2" { + t.Fatalf("image media_id = %q, want media-2", imageBody.Image.MediaID) + } + + streamBody, ok := commands[4].Body.(wecomRespondMsgBody) + if !ok { + t.Fatalf("unexpected finish body type %T", commands[4].Body) + } + if streamBody.MsgType != "stream" || streamBody.Stream == nil || !streamBody.Stream.Finish { + t.Fatalf("finish body = %+v", streamBody) + } + + if _, ok := ch.getTurn("chat-1"); ok { + t.Fatal("expected turn to be removed after media send") + } +} + +func TestSendMedia_SendsActiveFile(t *testing.T) { + t.Parallel() + + ch := newTestWeComChannel(t, bus.NewMessageBus()) + ch.SetRunning(true) + + store := media.NewFileMediaStore() + ch.SetMediaStore(store) + + filePath := filepath.Join(t.TempDir(), "report.pdf") + if err := os.WriteFile(filePath, []byte("%PDF-1.4"), 0o600); err != nil { + t.Fatalf("WriteFile() error = %v", err) + } + ref, err := store.Store(filePath, media.MediaMeta{ + Filename: "report.pdf", + ContentType: "application/pdf", + Source: "test", + CleanupPolicy: media.CleanupPolicyForgetOnly, + }, "scope-3") + if err != nil { + t.Fatalf("Store() error = %v", err) + } + + var commands []wecomCommand + ch.commandSend = func(cmd wecomCommand, _ time.Duration) (wecomEnvelope, error) { + commands = append(commands, cmd) + switch cmd.Cmd { + case wecomCmdUploadMediaInit: + return wecomTestAck(wecomUploadMediaInitResponse{UploadID: "upload-3"}), nil + case wecomCmdUploadMediaEnd: + return wecomTestAck(wecomUploadMediaFinishResponse{ + Type: "file", + MediaID: "media-3", + }), nil + default: + return wecomTestAck(nil), nil + } + } + + err = ch.SendMedia(context.Background(), bus.OutboundMediaMessage{ + Channel: "wecom", + ChatID: "chat-2", + Parts: []bus.MediaPart{{ + Ref: ref, + Type: "file", + Filename: "report.pdf", + ContentType: "application/pdf", + }}, + }) + if err != nil { + t.Fatalf("SendMedia() error = %v", err) + } + + if len(commands) != 4 { + t.Fatalf("expected 4 commands, got %d", len(commands)) + } + if commands[0].Cmd != wecomCmdUploadMediaInit { + t.Fatalf("first command = %q, want %q", commands[0].Cmd, wecomCmdUploadMediaInit) + } + initBody, ok := commands[0].Body.(wecomUploadMediaInitBody) + if !ok { + t.Fatalf("unexpected init body type %T", commands[0].Body) + } + if initBody.Type != "file" || initBody.Filename != "report.pdf" { + t.Fatalf("init body = %+v", initBody) + } + if commands[1].Cmd != wecomCmdUploadMediaChunk { + t.Fatalf("second command = %q, want %q", commands[1].Cmd, wecomCmdUploadMediaChunk) + } + if commands[2].Cmd != wecomCmdUploadMediaEnd { + t.Fatalf("third command = %q, want %q", commands[2].Cmd, wecomCmdUploadMediaEnd) + } + if commands[3].Cmd != wecomCmdSendMsg { + t.Fatalf("fourth command = %q, want %q", commands[3].Cmd, wecomCmdSendMsg) + } + + body, ok := commands[3].Body.(wecomSendMsgBody) + if !ok { + t.Fatalf("unexpected body type %T", commands[3].Body) + } + if body.MsgType != "file" || body.File == nil { + t.Fatalf("body = %+v", body) + } + if body.File.MediaID != "media-3" { + t.Fatalf("file media_id = %q, want media-3", body.File.MediaID) + } +} + +func newTestWeComChannel(t *testing.T, messageBus *bus.MessageBus) *WeComChannel { + t.Helper() + + cfg := config.WeComConfig{BotID: "bot-1"} + cfg.SetSecret("secret-1") + ch, err := NewChannel(cfg, messageBus) + if err != nil { + t.Fatalf("NewChannel() error = %v", err) + } + ch.ctx = context.Background() + ch.routes = newReqIDStore(filepath.Join(t.TempDir(), "reqids.json")) + return ch +} + +func wecomTestJPEGData(t *testing.T) []byte { + t.Helper() + + const jpegBase64 = "/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAP//////////////////////////////////////////////////////////////////////////////////////" + + "//////////////////////////////////////////////////////////////////////////////////////////////2wBDAf//////////////////////////////////////////////////////////////////////////////////////" + + "//////////////////////////////////////////////////////////////////////////////////////////////wAARCAABAAEDASIAAhEBAxEB/8QAFQABAQAAAAAAAAAAAAAAAAAAAAb/xAAVEQEBAAAAAAAAAAAAAAAAAAAABf/aAAwDAQACEAMQAAAB6A//xAAVEAEBAAAAAAAAAAAAAAAAAAAAEf/aAAgBAQABBQJf/8QAFBEBAAAAAAAAAAAAAAAAAAAAEP/aAAgBAwEBPwF//8QAFBEBAAAAAAAAAAAAAAAAAAAAEP/aAAgBAgEBPwF//8QAFBABAAAAAAAAAAAAAAAAAAAAEP/aAAgBAQAGPwJf/8QAFBABAAAAAAAAAAAAAAAAAAAAEP/aAAgBAQABPyFf/9k=" + + return decodeTestBase64(t, jpegBase64) +} + +func TestDecodeWeComUploadFinish_AcceptsNumericCreatedAt(t *testing.T) { + t.Parallel() + + resp, err := decodeWeComEnvelopeBody[wecomUploadMediaFinishResponse](wecomEnvelope{ + Body: json.RawMessage(`{"type":"file","media_id":"media-1","created_at":1380000000}`), + }) + if err != nil { + t.Fatalf("decodeWeComEnvelopeBody() error = %v", err) + } + if resp.Type != "file" || resp.MediaID != "media-1" { + t.Fatalf("response = %+v", resp) + } + if string(resp.CreatedAt) != "1380000000" { + t.Fatalf("created_at = %s, want 1380000000", string(resp.CreatedAt)) + } +} + +func wecomTestAck(body any) wecomEnvelope { + var raw []byte + if body != nil { + encoded, err := json.Marshal(body) + if err != nil { + panic(err) + } + raw = encoded + } + return wecomEnvelope{ + ErrCode: 0, + ErrMsg: "ok", + Body: raw, + } +} diff --git a/pkg/config/SECURITY_CONFIG.md b/pkg/config/SECURITY_CONFIG.md deleted file mode 100644 index c5aed54ae..000000000 --- a/pkg/config/SECURITY_CONFIG.md +++ /dev/null @@ -1,551 +0,0 @@ -# Security Configuration Refactoring - -## Overview - -This refactoring introduces a `.security.yml` file to store all sensitive data (API keys, tokens, secrets, passwords) separately from the main configuration. This improves security by: - -1. **Separation of concerns**: Configuration settings and secrets are in separate files -2. **Easier sharing**: The main config can be shared without exposing sensitive data -3. **Better version control**: `.security.yml` can be added to `.gitignore` -4. **Flexible deployment**: Different environments can use different security files - -## File Structure - -``` -~/.picoclaw/ -├── config.json # Main configuration (safe to share) -└── .security.yml # Security data (never share) -``` - -## Usage - -### Basic Configuration - -In your `config.json`, use `ref:` references to point to values in `.security.yml`: - -```json -{ - "version": 1, - "model_list": [ - { - "model_name": "gpt-5.4", - "model": "openai/gpt-5.4", - "api_base": "https://api.openai.com/v1", - "api_key": "ref:model_list.gpt-5.4.api_key" - } - ], - "channels": { - "telegram": { - "enabled": true, - "token": "ref:channels.telegram.token" - } - } -} -``` - -### Security Configuration - -In your `.security.yml`, store the actual values: - -```yaml -model_list: - gpt-5.4: - api_keys: - - "sk-your-actual-api-key-1" - - "sk-your-actual-api-key-2" # Optional: Multiple keys for failover - claude-sonnet-4.6: - api_keys: - - "sk-your-actual-anthropic-key" # Single key in array format - -channels: - telegram: - token: "your-telegram-bot-token" - -web: - brave: - api_keys: - - "BSAyour-brave-api-key-1" - - "BSAyour-brave-api-key-2" # Optional: Multiple keys for failover - tavily: - api_keys: - - "tvly-your-tavily-api-key" # Single key in array format - glm_search: - api_key: "your-glm-search-api-key" # GLMSearch uses single key format -``` - -## Reference Format - -### Model API Keys - -Format: `ref:model_list..api_key` - -Example: `ref:model_list.gpt-5.4.api_key` - -### Channel Tokens/Secrets - -Format: `ref:channels..` - -Examples: -- `ref:channels.telegram.token` -- `ref:channels.feishu.app_secret` -- `ref:channels.feishu.encrypt_key` -- `ref:channels.feishu.verification_token` -- `ref:channels.discord.token` -- `ref:channels.qq.app_secret` -- `ref:channels.dingtalk.client_secret` -- `ref:channels.slack.bot_token` -- `ref:channels.slack.app_token` -- `ref:channels.matrix.access_token` -- `ref:channels.line.channel_secret` -- `ref:channels.line.channel_access_token` -- `ref:channels.onebot.access_token` -- `ref:channels.wecom.token` -- `ref:channels.wecom.encoding_aes_key` -- `ref:channels.wecom_app.corp_secret` -- `ref:channels.wecom_app.token` -- `ref:channels.wecom_app.encoding_aes_key` -- `ref:channels.wecom_aibot.token` -- `ref:channels.wecom_aibot.encoding_aes_key` -- `ref:channels.pico.token` -- `ref:channels.irc.password` -- `ref:channels.irc.nickserv_password` -- `ref:channels.irc.sasl_password` - -### Web Tool API Keys - -Format: `ref:web..` - -Examples: -- `ref:web.brave.api_key` -- `ref:web.tavily.api_key` -- `ref:web.perplexity.api_key` -- `ref:web.glm_search.api_key` - -### Skills Registry Tokens - -Format: `ref:skills..` - -Examples: -- `ref:skills.github.token` -- `ref:skills.clawhub.auth_token` - -## Backward Compatibility - -The refactoring maintains full backward compatibility: - -1. **Direct values**: You can still use direct values in `config.json` (not recommended for production) -2. **Mixed usage**: You can mix `ref:` references and direct values -3. **Optional security file**: If `.security.yml` doesn't exist, all references will fail (but direct values still work) - -### API Key Formats in .security.yml - -**Models (gpt-5.4, claude-sonnet-4.6, etc.):** -- Must use `api_keys` (array) format -- Both single and multiple keys use array format - -**Web Tools (Brave, Tavily, Perplexity):** -- Must use `api_keys` (array) format -- Both single and multiple keys use array format - -**Web Tools (GLMSearch):** -- Must use `api_key` (single string) format -- Does NOT support array format - -**Channels (Telegram, Discord, etc.):** -- Use single field names (e.g., `token`, `app_secret`) -- Each channel uses its specific field names - -### Single Key (Models) - -Use array format with one element: -```yaml -model_list: - gpt-5.4: - api_keys: - - "sk-your-key" -``` - -In `config.json`: -```json -{ - "api_key": "ref:model_list.gpt-5.4.api_key" -} -``` - -### Single Key (GLMSearch) - -Use single string format: -```yaml -web: - glm_search: - api_key: "your-glm-key" -``` - -In `config.json`: -```json -{ - "api_key": "ref:web.glm_search.api_key" -} -``` - -## Migration Guide - -### Step 1: Create .security.yml - -Copy the example template: -```bash -cp security.example.yml ~/.picoclaw/.security.yml -``` - -### Step 2: Fill in your actual values - -Edit `~/.picoclaw/.security.yml` and replace placeholder values with your actual API keys and tokens. - -### Step 3: Update config.json - -Replace sensitive values in `~/.picoclaw/config.json` with `ref:` references: - -**Before:** -```json -{ - "model_list": [ - { - "model_name": "gpt-5.4", - "model": "openai/gpt-5.4", - "api_key": "sk-your-actual-api-key-here" - } - ] -} -``` - -**After:** -```json -{ - "model_list": [ - { - "model_name": "gpt-5.4", - "model": "openai/gpt-5.4", - "api_key": "ref:model_list.gpt-5.4.api_key" - } - ] -} -``` - -### Step 4: Verify - -Restart PicoClaw and verify it loads correctly: -```bash -picoclaw --version -``` - -## Security Best Practices - -1. **Never commit `.security.yml`** to version control -2. **Set file permissions**: `chmod 600 ~/.picoclaw/.security.yml` -3. **Use different keys** for different environments (dev, staging, production) -4. **Rotate keys regularly** and update `.security.yml` -5. **Backup securely**: Encrypt backups containing `.security.yml` - -## API - -### LoadSecurityConfig - -```go -func LoadSecurityConfig(securityPath string) (*SecurityConfig, error) -``` - -Loads the security configuration from `.security.yml`. Returns an empty `SecurityConfig` if the file doesn't exist. - -### SaveSecurityConfig - -```go -func SaveSecurityConfig(securityPath string, sec *SecurityConfig) error -``` - -Saves the security configuration to `.security.yml` with `0o600` permissions. - -### ResolveReference - -```go -func (sec *SecurityConfig) ResolveReference(ref string) (string, error) -``` - -Resolves a reference string (e.g., `"ref:model_list.test.api_key"`) and returns the actual value. - -### SecurityPath - -```go -func SecurityPath(configPath string) string -``` - -Returns the path to `.security.yml` relative to the config file. - -## Example: Complete Configuration - -### config.json -```json -{ - "version": 1, - "agents": { - "defaults": { - "workspace": "~/picoclaw-workspace", - "model_name": "gpt-5.4" - } - }, - "model_list": [ - { - "model_name": "gpt-5.4", - "model": "openai/gpt-5.4", - "api_base": "https://api.openai.com/v1", - "api_key": "ref:model_list.gpt-5.4.api_key" - }, - { - "model_name": "claude-sonnet-4.6", - "model": "anthropic/claude-sonnet-4.6", - "api_base": "https://api.anthropic.com/v1", - "api_key": "ref:model_list.claude-sonnet-4.6.api_key" - } - ], - "channels": { - "telegram": { - "enabled": true, - "token": "ref:channels.telegram.token" - } - }, - "tools": { - "web": { - "brave": { - "enabled": true, - "api_key": "ref:web.brave.api_key" - } - } - } -} -``` - -### .security.yml -```yaml -model_list: - gpt-5.4: - api_keys: - - "sk-proj-actual-openai-key-1" - - "sk-proj-actual-openai-key-2" - claude-sonnet-4.6: - api_keys: - - "sk-ant-actual-anthropic-key" # Single key in array format - -channels: - telegram: - token: "1234567890:ABCdefGHIjklMNOpqrsTUVwxyz" - -web: - brave: - api_keys: - - "BSAactualbravekey-1" - - "BSAactualbravekey-2" - tavily: - api_keys: - - "tvly-your-tavily-key" # Single key in array format - glm_search: - api_key: "your-glm-key" # GLMSearch uses single key format -``` - -## Testing - -The refactoring includes comprehensive tests: - -```bash -go test ./pkg/config -run TestSecurityConfig -``` - -## Troubleshooting - -### Error: "model security entry not found" - -- Ensure the model name in your reference matches exactly in `.security.yml` -- Check that the `model_list` section exists in `.security.yml` -- For models with indexed names (e.g., "gpt-5.4:0"), ensure the exact name is used or check the base name without index - -### Error: "failed to load security config" - -- Verify `.security.yml` exists in the same directory as `config.json` -- Check the YAML syntax is valid (use a YAML validator) -- Ensure file permissions allow reading - -### Error: "unknown reference path" - -- Verify the reference format is correct -- Check the path structure matches the examples above -- Ensure all required sections exist in `.security.yml` - -## Advanced Features - -### Multiple API Keys (Load Balancing & Failover) - -Both models and web tools support multiple API keys for improved reliability: - -**Benefits:** -- **Load balancing**: Requests are distributed across multiple keys -- **Failover**: Automatic switching to another key if one fails -- **Rate limit management**: Distribute usage across multiple keys -- **High availability**: Reduce downtime during API provider issues - -#### Example: Model with Multiple Keys - -**.security.yml:** -```yaml -model_list: - gpt-5.4: - api_keys: - - "sk-proj-key-1" - - "sk-proj-key-2" - - "sk-proj-key-3" -``` - -**config.json:** -```json -{ - "model_list": [ - { - "model_name": "gpt-5.4", - "model": "openai/gpt-5.4", - "api_key": "ref:model_list.gpt-5.4.api_key" - } - ] -} -``` - -#### Example: Web Tool with Multiple Keys - -**.security.yml:** -```yaml -web: - brave: - api_keys: - - "BSA-key-1" - - "BSA-key-2" - tavily: - api_keys: - - "tvly-your-key" # Single key in array format - glm_search: - api_key: "your-glm-key" # GLMSearch uses single key format -``` - -**config.json:** -```json -{ - "tools": { - "web": { - "brave": { - "enabled": true, - "api_key": "ref:web.brave.api_key" - }, - "tavily": { - "enabled": true, - "api_key": "ref:web.tavily.api_key" - } - } - } -} -``` - -#### Supported Formats - -**Models - Single key:** -```yaml -model_list: - gpt-5.4: - api_keys: - - "sk-your-key" # Array with one element -``` - -**Models - Multiple keys:** -```yaml -model_list: - gpt-5.4: - api_keys: - - "sk-your-key-1" - - "sk-your-key-2" - - "sk-your-key-3" -``` - -**Web Tools (Brave/Tavily/Perplexity) - Single key:** -```yaml -web: - brave: - api_keys: - - "BSA-your-key" # Array with one element -``` - -**Web Tools (Brave/Tavily/Perplexity) - Multiple keys:** -```yaml -web: - brave: - api_keys: - - "BSA-key-1" - - "BSA-key-2" -``` - -**Web Tool (GLMSearch) - Single key only:** -```yaml -web: - glm_search: - api_key: "your-glm-key" # Single string (NOT array) -``` - -All formats work identically in `config.json` - you always use the same reference format: -```json -{ - "api_key": "ref:model_list.gpt-5.4.api_key" -} -``` - -### Model Indexing for Load Balancing - -When you have multiple models with the same base name but different API keys, you can use indexed names: - -**.security.yml:** -```yaml -model_list: - gpt-5.4: - api_keys: - - "sk-proj-key-1" - - "sk-proj-key-2" -``` - -The system will automatically expand this into multiple model entries with fallback support. - -### Environment Variables - -You can override any security value using environment variables: - -**For models:** -```bash -export PICOCLAW_MODEL_LIST_GPT-5.4_API_KEY="sk-from-env" -``` - -**For channels:** -```bash -export PICOCLAW_CHANNELS_TELEGRAM_TOKEN="token-from-env" -``` - -**For web tools:** -```bash -export PICOCLAW_WEB_BRAVE_API_KEY="key-from-env" -``` - -Environment variables follow this pattern: `PICOCLAW_
___` with dots replaced by underscores and converted to uppercase. - -### Multiple API Keys Not Working - -- Ensure you're using `api_keys` (plural) in `.security.yml` for models and web tools (except GLMSearch) -- Check that the array format is correct in YAML (proper indentation) -- Remember: Models, Brave, Tavily, Perplexity MUST use `api_keys` (array format) -- GLMSearch MUST use `api_key` (single string format) -- The reference in `config.json` is the same regardless of single or multiple keys - -### Load Balancing/Failover Issues - -- Verify all API keys in the `api_keys` array are valid -- Check that all keys have the same rate limits and permissions -- Monitor logs to see which keys are being used and failing diff --git a/pkg/config/config.go b/pkg/config/config.go index 84e1ab61a..2780fb401 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -3,6 +3,7 @@ package config import ( "encoding/json" "fmt" + "math/rand" "os" "path/filepath" "strings" @@ -26,6 +27,22 @@ var rrCounter atomic.Uint64 type FlexibleStringSlice []string func (f *FlexibleStringSlice) UnmarshalJSON(data []byte) error { + // Accept a single JSON string for convenience, e.g.: + // "text": "Thinking..." + var singleString string + if err := json.Unmarshal(data, &singleString); err == nil { + *f = FlexibleStringSlice{singleString} + return nil + } + + // Accept a single JSON number too, to keep symmetry with mixed allow_from + // payloads that may contain numeric identifiers. + var singleNumber float64 + if err := json.Unmarshal(data, &singleNumber); err == nil { + *f = FlexibleStringSlice{fmt.Sprintf("%.0f", singleNumber)} + return nil + } + // Try []string first var ss []string if err := json.Unmarshal(data, &ss); err == nil { @@ -319,12 +336,10 @@ type AgentDefaults struct { SteeringMode string `json:"steering_mode,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_STEERING_MODE"` // "one-at-a-time" (default) or "all" SubTurn SubTurnConfig `json:"subturn" envPrefix:"PICOCLAW_AGENTS_DEFAULTS_SUBTURN_"` ToolFeedback ToolFeedbackConfig `json:"tool_feedback,omitempty"` + SplitOnMarker bool `json:"split_on_marker" env:"PICOCLAW_AGENTS_DEFAULTS_SPLIT_ON_MARKER"` // split messages on <|[SPLIT]|> marker } -const ( - DefaultMaxMediaSize = 20 * 1024 * 1024 // 20 MB - DefaultWeComAIBotProcessingMessage = "⏳ Processing, please wait. The results will be sent shortly." -) +const DefaultMaxMediaSize = 20 * 1024 * 1024 // 20 MB func (d *AgentDefaults) GetMaxMediaSize() int { if d.MaxMediaSize > 0 { @@ -364,9 +379,7 @@ type ChannelsConfig struct { Matrix MatrixConfig `json:"matrix"` LINE LINEConfig `json:"line"` OneBot OneBotConfig `json:"onebot"` - WeCom WeComConfig `json:"wecom"` - WeComApp WeComAppConfig `json:"wecom_app"` - WeComAIBot WeComAIBotConfig `json:"wecom_aibot"` + WeCom WeComConfig `json:"wecom" envPrefix:"PICOCLAW_CHANNELS_WECOM_"` Weixin WeixinConfig `json:"weixin"` Pico PicoConfig `json:"pico"` PicoClient PicoClientConfig `json:"pico_client"` @@ -386,8 +399,20 @@ type TypingConfig struct { // PlaceholderConfig controls placeholder message behavior (Phase 10). type PlaceholderConfig struct { - Enabled bool `json:"enabled,omitempty"` - Text string `json:"text,omitempty"` + Enabled bool `json:"enabled"` + Text FlexibleStringSlice `json:"text,omitempty"` +} + +// GetRandomText returns a random placeholder text, or default if none set. +func (p *PlaceholderConfig) GetRandomText() string { + if len(p.Text) == 0 { + return "Thinking..." + } + if len(p.Text) == 1 { + return p.Text[0] + } + idx := rand.Intn(len(p.Text)) + return p.Text[idx] } type StreamingConfig struct { @@ -591,18 +616,20 @@ func (c *SlackConfig) SetAppToken(token string) { } type MatrixConfig struct { - Enabled bool `json:"enabled" env:"PICOCLAW_CHANNELS_MATRIX_ENABLED"` - Homeserver string `json:"homeserver" env:"PICOCLAW_CHANNELS_MATRIX_HOMESERVER"` - UserID string `json:"user_id" env:"PICOCLAW_CHANNELS_MATRIX_USER_ID"` + Enabled bool `json:"enabled" env:"PICOCLAW_CHANNELS_MATRIX_ENABLED"` + Homeserver string `json:"homeserver" env:"PICOCLAW_CHANNELS_MATRIX_HOMESERVER"` + UserID string `json:"user_id" env:"PICOCLAW_CHANNELS_MATRIX_USER_ID"` accessToken string - DeviceID string `json:"device_id,omitempty" env:"PICOCLAW_CHANNELS_MATRIX_DEVICE_ID"` - JoinOnInvite bool `json:"join_on_invite" env:"PICOCLAW_CHANNELS_MATRIX_JOIN_ON_INVITE"` - MessageFormat string `json:"message_format,omitempty" env:"PICOCLAW_CHANNELS_MATRIX_MESSAGE_FORMAT"` - AllowFrom FlexibleStringSlice `json:"allow_from" env:"PICOCLAW_CHANNELS_MATRIX_ALLOW_FROM"` + DeviceID string `json:"device_id,omitempty" env:"PICOCLAW_CHANNELS_MATRIX_DEVICE_ID"` + JoinOnInvite bool `json:"join_on_invite" env:"PICOCLAW_CHANNELS_MATRIX_JOIN_ON_INVITE"` + MessageFormat string `json:"message_format,omitempty" env:"PICOCLAW_CHANNELS_MATRIX_MESSAGE_FORMAT"` + AllowFrom FlexibleStringSlice `json:"allow_from" env:"PICOCLAW_CHANNELS_MATRIX_ALLOW_FROM"` GroupTrigger GroupTriggerConfig `json:"group_trigger,omitempty"` Placeholder PlaceholderConfig `json:"placeholder,omitempty"` - ReasoningChannelID string `json:"reasoning_channel_id" env:"PICOCLAW_CHANNELS_MATRIX_REASONING_CHANNEL_ID"` + ReasoningChannelID string `json:"reasoning_channel_id" env:"PICOCLAW_CHANNELS_MATRIX_REASONING_CHANNEL_ID"` secDirty bool + CryptoDatabasePath string `json:"crypto_database_path,omitempty" env:"PICOCLAW_CHANNELS_MATRIX_CRYPTO_DATABASE_PATH"` + CryptoPassphrase string `json:"crypto_passphrase,omitempty" env:"PICOCLAW_CHANNELS_MATRIX_CRYPTO_PASSPHRASE"` } // AccessToken returns the Matrix access token @@ -678,136 +705,28 @@ func (c *OneBotConfig) SetAccessToken(token string) { c.secDirty = true } +type WeComGroupConfig struct { + AllowFrom FlexibleStringSlice `json:"allow_from,omitempty"` +} + type WeComConfig struct { - Enabled bool `json:"enabled" env:"PICOCLAW_CHANNELS_WECOM_ENABLED"` - token string - encodingAESKey string - WebhookURL string `json:"webhook_url" env:"PICOCLAW_CHANNELS_WECOM_WEBHOOK_URL"` - WebhookHost string `json:"webhook_host" env:"PICOCLAW_CHANNELS_WECOM_WEBHOOK_HOST"` - WebhookPort int `json:"webhook_port" env:"PICOCLAW_CHANNELS_WECOM_WEBHOOK_PORT"` - WebhookPath string `json:"webhook_path" env:"PICOCLAW_CHANNELS_WECOM_WEBHOOK_PATH"` - AllowFrom FlexibleStringSlice `json:"allow_from" env:"PICOCLAW_CHANNELS_WECOM_ALLOW_FROM"` - ReplyTimeout int `json:"reply_timeout" env:"PICOCLAW_CHANNELS_WECOM_REPLY_TIMEOUT"` - GroupTrigger GroupTriggerConfig `json:"group_trigger,omitempty"` - ReasoningChannelID string `json:"reasoning_channel_id" env:"PICOCLAW_CHANNELS_WECOM_REASONING_CHANNEL_ID"` - secDirty bool + Enabled bool `json:"enabled" env:"ENABLED"` + BotID string `json:"bot_id" env:"BOT_ID"` + secret string + WebSocketURL string `json:"websocket_url,omitempty" env:"WEBSOCKET_URL"` + SendThinkingMessage bool `json:"send_thinking_message" env:"SEND_THINKING_MESSAGE"` + AllowFrom FlexibleStringSlice `json:"allow_from" env:"ALLOW_FROM"` + ReasoningChannelID string `json:"reasoning_channel_id" env:"REASONING_CHANNEL_ID"` + secDirty bool } -// Token returns the WeCom token -func (c *WeComConfig) Token() string { - return c.token -} - -// SetToken sets the WeCom token -func (c *WeComConfig) SetToken(token string) { - c.token = token - c.secDirty = true -} - -// EncodingAESKey returns the WeCom encoding AES key -func (c *WeComConfig) EncodingAESKey() string { - return c.encodingAESKey -} - -// SetEncodingAESKey sets the WeCom encoding AES key -func (c *WeComConfig) SetEncodingAESKey(key string) { - c.encodingAESKey = key - c.secDirty = true -} - -type WeComAppConfig struct { - Enabled bool `json:"enabled" env:"PICOCLAW_CHANNELS_WECOM_APP_ENABLED"` - CorpID string `json:"corp_id" env:"PICOCLAW_CHANNELS_WECOM_APP_CORP_ID"` - corpSecret string - AgentID int64 `json:"agent_id" env:"PICOCLAW_CHANNELS_WECOM_APP_AGENT_ID"` - token string - encodingAESKey string - WebhookHost string `json:"webhook_host" env:"PICOCLAW_CHANNELS_WECOM_APP_WEBHOOK_HOST"` - WebhookPort int `json:"webhook_port" env:"PICOCLAW_CHANNELS_WECOM_APP_WEBHOOK_PORT"` - WebhookPath string `json:"webhook_path" env:"PICOCLAW_CHANNELS_WECOM_APP_WEBHOOK_PATH"` - AllowFrom FlexibleStringSlice `json:"allow_from" env:"PICOCLAW_CHANNELS_WECOM_APP_ALLOW_FROM"` - ReplyTimeout int `json:"reply_timeout" env:"PICOCLAW_CHANNELS_WECOM_APP_REPLY_TIMEOUT"` - GroupTrigger GroupTriggerConfig `json:"group_trigger,omitempty"` - ReasoningChannelID string `json:"reasoning_channel_id" env:"PICOCLAW_CHANNELS_WECOM_APP_REASONING_CHANNEL_ID"` - secDirty bool -} - -// CorpSecret returns the corporate secret for WeCom app -func (c *WeComAppConfig) CorpSecret() string { - return c.corpSecret -} - -// SetCorpSecret sets the corporate secret for WeCom app -func (c *WeComAppConfig) SetCorpSecret(secret string) { - c.corpSecret = secret - c.secDirty = true -} - -// Token returns the webhook token for WeCom app -func (c *WeComAppConfig) Token() string { - return c.token -} - -// SetToken sets the webhook token for WeCom app -func (c *WeComAppConfig) SetToken(token string) { - c.token = token - c.secDirty = true -} - -// EncodingAESKey returns the encoding AES key for WeCom app -func (c *WeComAppConfig) EncodingAESKey() string { - return c.encodingAESKey -} - -// SetEncodingAESKey sets the encoding AES key for WeCom app -func (c *WeComAppConfig) SetEncodingAESKey(key string) { - c.encodingAESKey = key - c.secDirty = true -} - -type WeComAIBotConfig struct { - Enabled bool `json:"enabled" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_ENABLED"` - BotID string `json:"bot_id,omitempty" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_BOT_ID"` - secret string - token string - encodingAESKey string - WebhookPath string `json:"webhook_path,omitempty" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_WEBHOOK_PATH"` - AllowFrom FlexibleStringSlice `json:"allow_from" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_ALLOW_FROM"` - ReplyTimeout int `json:"reply_timeout" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_REPLY_TIMEOUT"` - MaxSteps int `json:"max_steps" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_MAX_STEPS"` // Maximum streaming steps - WelcomeMessage string `json:"welcome_message" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_WELCOME_MESSAGE"` // Sent on enter_chat event; empty = no welcome - ProcessingMessage string `json:"processing_message,omitempty" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_PROCESSING_MESSAGE"` - ReasoningChannelID string `json:"reasoning_channel_id" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_REASONING_CHANNEL_ID"` - secDirty bool -} - -// Token returns the webhook token for WeCom AI bot -func (c *WeComAIBotConfig) Token() string { - return c.token -} - -// EncodingAESKey returns the encoding AES key for WeCom AI bot -func (c *WeComAIBotConfig) EncodingAESKey() string { - return c.encodingAESKey -} - -// SetToken sets the token for WeCom AI bot -func (c *WeComAIBotConfig) SetToken(token string) { - c.token = token - c.secDirty = true -} - -// SetEncodingAESKey sets the encoding AES key for WeCom AI bot -func (c *WeComAIBotConfig) SetEncodingAESKey(key string) { - c.encodingAESKey = key - c.secDirty = true -} - -func (c *WeComAIBotConfig) Secret() string { +// Secret returns the WeCom bot secret. +func (c *WeComConfig) Secret() string { return c.secret } -func (c *WeComAIBotConfig) SetSecret(secret string) { +// SetSecret sets the WeCom bot secret. +func (c *WeComConfig) SetSecret(secret string) { c.secret = secret c.secDirty = true } @@ -968,6 +887,10 @@ type ModelConfig struct { secModelName string apiKeys []string secDirty bool + + // isVirtual marks this model as a virtual model generated from multi-key expansion. + // Virtual models should not be persisted to config files. + isVirtual bool } // APIKey returns the first API key from apiKeys @@ -978,6 +901,11 @@ func (c *ModelConfig) APIKey() string { return "" } +// IsVirtual returns true if this model was generated from multi-key expansion. +func (c *ModelConfig) IsVirtual() bool { + return c.isVirtual +} + // Validate checks if the ModelConfig has all required fields. func (c *ModelConfig) Validate() error { if c.ModelName == "" { @@ -1416,17 +1344,29 @@ func LoadConfig(path string) (*Config, error) { if err != nil { return nil, err } - // Load security configuration - securityPath := securityPath(path) - sec, err := loadSecurityConfig(securityPath) + + // Legacy config (no version field) + tmpCfg, e := loadConfigV0(data) + if e != nil { + return nil, e + } + + tmpCfgMigrated, e := tmpCfg.Migrate() + if e != nil { + logger.ErrorF("config migrate fail", map[string]any{"from": versionInfo.Version, "to": CurrentVersion}) + return nil, e + } + + // Load security configuration from .security.yml + secPath := securityPath(path) + sec, err := loadSecurityConfig(secPath) if err != nil { return nil, fmt.Errorf("failed to load security config: %w", err) } - // Apply security references from .security.yml BEFORE resolveAPIKeys - // This resolves ref: references to actual values - if err := applySecurityConfig(cfg, sec); err != nil { - return nil, fmt.Errorf("failed to apply security config: %w", err) + // Merge security configs: config.json takes precedence over .security.yml + if err := applySecurityConfigWithPrecedence(cfg, tmpCfgMigrated, sec); err != nil { + return nil, fmt.Errorf("failed to merge security config: %w", err) } default: return nil, fmt.Errorf("unsupported config version: %d", versionInfo.Version) @@ -1623,39 +1563,10 @@ func applySecurityConfig(cfg *Config, sec *SecurityConfig) error { cfg.Channels.OneBot.accessToken = sec.Channels.OneBot.AccessToken } - // Handle WeCom token and encoding key + // Handle WeCom bot secret if sec.Channels.WeCom != nil { - if sec.Channels.WeCom.Token != "" { - cfg.Channels.WeCom.token = sec.Channels.WeCom.Token - } - if sec.Channels.WeCom.EncodingAESKey != "" { - cfg.Channels.WeCom.encodingAESKey = sec.Channels.WeCom.EncodingAESKey - } - } - - // Handle WeCom App credentials - if sec.Channels.WeComApp != nil { - if sec.Channels.WeComApp.CorpSecret != "" { - cfg.Channels.WeComApp.corpSecret = sec.Channels.WeComApp.CorpSecret - } - if sec.Channels.WeComApp.Token != "" { - cfg.Channels.WeComApp.token = sec.Channels.WeComApp.Token - } - if sec.Channels.WeComApp.EncodingAESKey != "" { - cfg.Channels.WeComApp.encodingAESKey = sec.Channels.WeComApp.EncodingAESKey - } - } - - // Handle WeCom AI Bot credentials - if sec.Channels.WeComAIBot != nil { - if sec.Channels.WeComAIBot.Token != "" { - cfg.Channels.WeComAIBot.token = sec.Channels.WeComAIBot.Token - } - if sec.Channels.WeComAIBot.EncodingAESKey != "" { - cfg.Channels.WeComAIBot.encodingAESKey = sec.Channels.WeComAIBot.EncodingAESKey - } - if sec.Channels.WeComAIBot.Secret != "" { - cfg.Channels.WeComAIBot.secret = sec.Channels.WeComAIBot.Secret + if sec.Channels.WeCom.Secret != "" { + cfg.Channels.WeCom.secret = sec.Channels.WeCom.Secret } } @@ -1688,6 +1599,28 @@ func applySecurityConfig(cfg *Config, sec *SecurityConfig) error { return nil } +// applySecurityConfigWithPrecedence merges security config from tmpCfg (migrated from configV0) and sec (SecurityConfig), +// with tmpCfg taking precedence. It then applies the merged security config to cfg. +func applySecurityConfigWithPrecedence(cfg *Config, tmpCfg *Config, sec *SecurityConfig) error { + // Get security config from tmpCfg (already extracted during migration) + var tmpSec *SecurityConfig + if tmpCfg != nil { + tmpSec = tmpCfg.security + } + + // If tmpCfg has no security config, just apply sec directly + if tmpSec == nil { + return applySecurityConfig(cfg, sec) + } + + // Merge sec and tmpSec, with tmpSec (from config.json) taking precedence + // mergeSecurityConfig(existing, newer) - newer takes precedence + mergedSec := mergeSecurityConfig(sec, tmpSec) + + // Apply the merged security config to cfg + return applySecurityConfig(cfg, mergedSec) +} + func toNameIndex(list []*ModelConfig) []string { nameList := make([]string, 0, len(list)) countMap := make(map[string]int) @@ -1879,27 +1812,10 @@ func SaveConfig(path string, cfg *Config) error { } if cfg.Channels.WeCom.secDirty { cfg.security.Channels.WeCom = &WeComSecurity{ - Token: cfg.Channels.WeCom.Token(), - EncodingAESKey: cfg.Channels.WeCom.EncodingAESKey(), + Secret: cfg.Channels.WeCom.Secret(), } cfg.Channels.WeCom.secDirty = false } - if cfg.Channels.WeComApp.secDirty { - cfg.security.Channels.WeComApp = &WeComAppSecurity{ - CorpSecret: cfg.Channels.WeComApp.CorpSecret(), - Token: cfg.Channels.WeComApp.Token(), - EncodingAESKey: cfg.Channels.WeComApp.EncodingAESKey(), - } - cfg.Channels.WeComApp.secDirty = false - } - if cfg.Channels.WeComAIBot.secDirty { - cfg.security.Channels.WeComAIBot = &WeComAIBotSecurity{ - Token: cfg.Channels.WeComAIBot.Token(), - EncodingAESKey: cfg.Channels.WeComAIBot.EncodingAESKey(), - Secret: cfg.Channels.WeComAIBot.Secret(), - } - cfg.Channels.WeComAIBot.secDirty = false - } if cfg.Tools.Web.Brave.secDirty { cfg.security.Web.Brave = &BraveSecurity{ APIKeys: cfg.Tools.Web.Brave.APIKeys(), @@ -1957,7 +1873,20 @@ func SaveConfig(path string, cfg *Config) error { return err } + // Filter out virtual models before serializing to config file + nonVirtualModels := make([]*ModelConfig, 0, len(cfg.ModelList)) + for _, m := range cfg.ModelList { + if !m.isVirtual { + nonVirtualModels = append(nonVirtualModels, m) + } + } + // Temporarily replace ModelList with filtered version for serialization + originalModelList := cfg.ModelList + cfg.ModelList = nonVirtualModels + data, err := json.MarshalIndent(cfg, "", " ") + // Restore original ModelList after serialization + cfg.ModelList = originalModelList if err != nil { return err } @@ -2175,6 +2104,7 @@ func expandMultiKeyModels(models []*ModelConfig) []*ModelConfig { RequestTimeout: m.RequestTimeout, ThinkingLevel: m.ThinkingLevel, ExtraBody: m.ExtraBody, + isVirtual: true, } expanded = append(expanded, additionalEntry) fallbackNames = append(fallbackNames, expandedName) diff --git a/pkg/config/config_old.go b/pkg/config/config_old.go index 01909f5a9..ad31833a3 100644 --- a/pkg/config/config_old.go +++ b/pkg/config/config_old.go @@ -85,23 +85,21 @@ type toolsConfigV0 struct { } type channelsConfigV0 struct { - WhatsApp WhatsAppConfig `json:"whatsapp"` - Telegram telegramConfigV0 `json:"telegram"` - Feishu feishuConfigV0 `json:"feishu"` - Discord discordConfigV0 `json:"discord"` - MaixCam maixcamConfigV0 `json:"maixcam"` - Weixin weixinConfigV0 `json:"weixin"` - QQ qqConfigV0 `json:"qq"` - DingTalk dingtalkConfigV0 `json:"dingtalk"` - Slack slackConfigV0 `json:"slack"` - Matrix matrixConfigV0 `json:"matrix"` - LINE lineConfigV0 `json:"line"` - OneBot onebotConfigV0 `json:"onebot"` - WeCom wecomConfigV0 `json:"wecom"` - WeComApp wecomappConfigV0 `json:"wecom_app"` - WeComAIBot wecomaibotConfigV0 `json:"wecom_aibot"` - Pico picoConfigV0 `json:"pico"` - IRC ircConfigV0 `json:"irc"` + WhatsApp WhatsAppConfig `json:"whatsapp"` + Telegram telegramConfigV0 `json:"telegram"` + Feishu feishuConfigV0 `json:"feishu"` + Discord discordConfigV0 `json:"discord"` + MaixCam maixcamConfigV0 `json:"maixcam"` + Weixin weixinConfigV0 `json:"weixin"` + QQ qqConfigV0 `json:"qq"` + DingTalk dingtalkConfigV0 `json:"dingtalk"` + Slack slackConfigV0 `json:"slack"` + Matrix matrixConfigV0 `json:"matrix"` + LINE lineConfigV0 `json:"line"` + OneBot onebotConfigV0 `json:"onebot"` + WeCom wecomConfigV0 `json:"wecom" envPrefix:"PICOCLAW_CHANNELS_WECOM_"` + Pico picoConfigV0 `json:"pico"` + IRC ircConfigV0 `json:"irc"` } func (v *channelsConfigV0) ToChannelsConfig() (ChannelsConfig, ChannelsSecurity) { @@ -117,45 +115,39 @@ func (v *channelsConfigV0) ToChannelsConfig() (ChannelsConfig, ChannelsSecurity) line, lineSecurity := v.LINE.ToLINEConfig() onebot, onebotSecurity := v.OneBot.ToOneBotConfig() wecom, wecomSecurity := v.WeCom.ToWeComConfig() - wecomapp, wecomappSecurity := v.WeComApp.ToWeComAppConfig() - wecomaibot, wecomaibotSecurity := v.WeComAIBot.ToWeComAIBotConfig() pico, picoSecurity := v.Pico.ToPicoConfig() irc, ircSecurity := v.IRC.ToIRCConfig() return ChannelsConfig{ - WhatsApp: v.WhatsApp, - Telegram: telegram, - Feishu: feishu, - Discord: discord, - MaixCam: maixcam, - QQ: qq, - Weixin: weixin, - DingTalk: dingtalk, - Slack: slack, - Matrix: matrix, - LINE: line, - OneBot: onebot, - WeCom: wecom, - WeComApp: wecomapp, - WeComAIBot: wecomaibot, - Pico: pico, - IRC: irc, + WhatsApp: v.WhatsApp, + Telegram: telegram, + Feishu: feishu, + Discord: discord, + MaixCam: maixcam, + QQ: qq, + Weixin: weixin, + DingTalk: dingtalk, + Slack: slack, + Matrix: matrix, + LINE: line, + OneBot: onebot, + WeCom: wecom, + Pico: pico, + IRC: irc, }, ChannelsSecurity{ - Telegram: telegramSecurity, - Feishu: feishuSecurity, - Discord: discordSecurity, - QQ: qqSecurity, - Weixin: weixinSecurity, - DingTalk: dingtalkSecurity, - Slack: slackSecurity, - Matrix: matrixSecurity, - LINE: lineSecurity, - OneBot: onebotSecurity, - WeCom: wecomSecurity, - WeComApp: wecomappSecurity, - WeComAIBot: wecomaibotSecurity, - Pico: picoSecurity, - IRC: ircSecurity, + Telegram: telegramSecurity, + Feishu: feishuSecurity, + Discord: discordSecurity, + QQ: qqSecurity, + Weixin: weixinSecurity, + DingTalk: dingtalkSecurity, + Slack: slackSecurity, + Matrix: matrixSecurity, + LINE: lineSecurity, + OneBot: onebotSecurity, + WeCom: wecomSecurity, + Pico: picoSecurity, + IRC: ircSecurity, } } @@ -473,39 +465,32 @@ func (v *onebotConfigV0) ToOneBotConfig() (OneBotConfig, *OneBotSecurity) { } type wecomConfigV0 struct { - Enabled bool `json:"enabled" env:"PICOCLAW_CHANNELS_WECOM_ENABLED"` - Token string `json:"token" env:"PICOCLAW_CHANNELS_WECOM_TOKEN"` - EncodingAESKey string `json:"encoding_aes_key" env:"PICOCLAW_CHANNELS_WECOM_ENCODING_AES_KEY"` - WebhookURL string `json:"webhook_url" env:"PICOCLAW_CHANNELS_WECOM_WEBHOOK_URL"` - WebhookHost string `json:"webhook_host" env:"PICOCLAW_CHANNELS_WECOM_WEBHOOK_HOST"` - WebhookPort int `json:"webhook_port" env:"PICOCLAW_CHANNELS_WECOM_WEBHOOK_PORT"` - WebhookPath string `json:"webhook_path" env:"PICOCLAW_CHANNELS_WECOM_WEBHOOK_PATH"` - AllowFrom FlexibleStringSlice `json:"allow_from" env:"PICOCLAW_CHANNELS_WECOM_ALLOW_FROM"` - ReplyTimeout int `json:"reply_timeout" env:"PICOCLAW_CHANNELS_WECOM_REPLY_TIMEOUT"` - GroupTrigger GroupTriggerConfig `json:"group_trigger,omitempty"` - ReasoningChannelID string `json:"reasoning_channel_id" env:"PICOCLAW_CHANNELS_WECOM_REASONING_CHANNEL_ID"` + Enabled bool `json:"enabled" env:"ENABLED"` + BotID string `json:"bot_id" env:"BOT_ID"` + Secret string `json:"secret" env:"SECRET"` + WebSocketURL string `json:"websocket_url,omitempty" env:"WEBSOCKET_URL"` + SendThinkingMessage bool `json:"send_thinking_message" env:"SEND_THINKING_MESSAGE"` + DMPolicy string `json:"dm_policy,omitempty" env:"DM_POLICY"` + AllowFrom FlexibleStringSlice `json:"allow_from" env:"ALLOW_FROM"` + GroupPolicy string `json:"group_policy,omitempty" env:"GROUP_POLICY"` + GroupAllowFrom FlexibleStringSlice `json:"group_allow_from,omitempty" env:"GROUP_ALLOW_FROM"` + Groups map[string]WeComGroupConfig `json:"groups,omitempty"` + ReasoningChannelID string `json:"reasoning_channel_id" env:"REASONING_CHANNEL_ID"` } func (v *wecomConfigV0) ToWeComConfig() (WeComConfig, *WeComSecurity) { var sec *WeComSecurity - if v.Token != "" || v.EncodingAESKey != "" { - sec = &WeComSecurity{ - Token: v.Token, - EncodingAESKey: v.EncodingAESKey, - } + if v.Secret != "" { + sec = &WeComSecurity{Secret: v.Secret} } return WeComConfig{ - Enabled: v.Enabled, - token: v.Token, - encodingAESKey: v.EncodingAESKey, - WebhookURL: v.WebhookURL, - WebhookHost: v.WebhookHost, - WebhookPort: v.WebhookPort, - WebhookPath: v.WebhookPath, - AllowFrom: v.AllowFrom, - ReplyTimeout: v.ReplyTimeout, - GroupTrigger: v.GroupTrigger, - ReasoningChannelID: v.ReasoningChannelID, + Enabled: v.Enabled, + BotID: v.BotID, + secret: v.Secret, + WebSocketURL: v.WebSocketURL, + SendThinkingMessage: v.SendThinkingMessage, + AllowFrom: v.AllowFrom, + ReasoningChannelID: v.ReasoningChannelID, }, sec } @@ -537,81 +522,6 @@ func (v *weixinConfigV0) ToWeiXinConfig() (WeixinConfig, *WeixinSecurity) { }, sec } -type wecomappConfigV0 struct { - Enabled bool `json:"enabled" env:"PICOCLAW_CHANNELS_WECOM_APP_ENABLED"` - CorpID string `json:"corp_id" env:"PICOCLAW_CHANNELS_WECOM_APP_CORP_ID"` - CorpSecret string `json:"corp_secret" env:"PICOCLAW_CHANNELS_WECOM_APP_CORP_SECRET"` - AgentID int64 `json:"agent_id" env:"PICOCLAW_CHANNELS_WECOM_APP_AGENT_ID"` - Token string `json:"token" env:"PICOCLAW_CHANNELS_WECOM_APP_TOKEN"` - EncodingAESKey string `json:"encoding_aes_key" env:"PICOCLAW_CHANNELS_WECOM_APP_ENCODING_AES_KEY"` - WebhookHost string `json:"webhook_host" env:"PICOCLAW_CHANNELS_WECOM_APP_WEBHOOK_HOST"` - WebhookPort int `json:"webhook_port" env:"PICOCLAW_CHANNELS_WECOM_APP_WEBHOOK_PORT"` - WebhookPath string `json:"webhook_path" env:"PICOCLAW_CHANNELS_WECOM_APP_WEBHOOK_PATH"` - AllowFrom FlexibleStringSlice `json:"allow_from" env:"PICOCLAW_CHANNELS_WECOM_APP_ALLOW_FROM"` - ReplyTimeout int `json:"reply_timeout" env:"PICOCLAW_CHANNELS_WECOM_APP_REPLY_TIMEOUT"` - GroupTrigger GroupTriggerConfig `json:"group_trigger,omitempty"` - ReasoningChannelID string `json:"reasoning_channel_id" env:"PICOCLAW_CHANNELS_WECOM_APP_REASONING_CHANNEL_ID"` -} - -func (v *wecomappConfigV0) ToWeComAppConfig() (WeComAppConfig, *WeComAppSecurity) { - var sec *WeComAppSecurity - if v.CorpSecret != "" || v.Token != "" || v.EncodingAESKey != "" { - sec = &WeComAppSecurity{ - CorpSecret: v.CorpSecret, - Token: v.Token, - EncodingAESKey: v.EncodingAESKey, - } - } - return WeComAppConfig{ - Enabled: v.Enabled, - CorpID: v.CorpID, - corpSecret: v.CorpSecret, - AgentID: v.AgentID, - token: v.Token, - encodingAESKey: v.EncodingAESKey, - WebhookHost: v.WebhookHost, - WebhookPort: v.WebhookPort, - WebhookPath: v.WebhookPath, - AllowFrom: v.AllowFrom, - ReplyTimeout: v.ReplyTimeout, - GroupTrigger: v.GroupTrigger, - ReasoningChannelID: v.ReasoningChannelID, - }, sec -} - -type wecomaibotConfigV0 struct { - Enabled bool `json:"enabled" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_ENABLED"` - Token string `json:"token" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_TOKEN"` - Secret string `json:"secret" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_SECRET"` - EncodingAESKey string `json:"encoding_aes_key" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_ENCODING_AES_KEY"` - WebhookPath string `json:"webhook_path" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_WEBHOOK_PATH"` - AllowFrom FlexibleStringSlice `json:"allow_from" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_ALLOW_FROM"` - ReplyTimeout int `json:"reply_timeout" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_REPLY_TIMEOUT"` - MaxSteps int `json:"max_steps" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_MAX_STEPS"` - WelcomeMessage string `json:"welcome_message" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_WELCOME_MESSAGE"` - ReasoningChannelID string `json:"reasoning_channel_id" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_REASONING_CHANNEL_ID"` -} - -func (v *wecomaibotConfigV0) ToWeComAIBotConfig() (WeComAIBotConfig, *WeComAIBotSecurity) { - var sec *WeComAIBotSecurity - if v.Token != "" || v.Secret != "" || v.EncodingAESKey != "" { - sec = &WeComAIBotSecurity{ - Token: v.Token, - Secret: v.Secret, - EncodingAESKey: v.EncodingAESKey, - } - } - return WeComAIBotConfig{ - Enabled: v.Enabled, - WebhookPath: v.WebhookPath, - AllowFrom: v.AllowFrom, - ReplyTimeout: v.ReplyTimeout, - MaxSteps: v.MaxSteps, - WelcomeMessage: v.WelcomeMessage, - ReasoningChannelID: v.ReasoningChannelID, - }, sec -} - type picoConfigV0 struct { Enabled bool `json:"enabled" env:"PICOCLAW_CHANNELS_PICO_ENABLED"` Token string `json:"token" env:"PICOCLAW_CHANNELS_PICO_TOKEN"` @@ -923,6 +833,7 @@ type webToolsConfigV0 struct { Perplexity perplexityConfigV0 ` json:"perplexity"` SearXNG SearXNGConfig ` json:"searxng"` GLMSearch glmSearchConfigV0 ` json:"glm_search"` + BaiduSearch baiduSearchConfigV0 ` json:"baidu_search"` PreferNative bool ` json:"prefer_native" env:"PICOCLAW_TOOLS_WEB_PREFER_NATIVE"` Proxy string ` json:"proxy,omitempty" env:"PICOCLAW_TOOLS_WEB_PROXY"` FetchLimitBytes int64 ` json:"fetch_limit_bytes,omitempty" env:"PICOCLAW_TOOLS_WEB_FETCH_LIMIT_BYTES"` @@ -1014,11 +925,34 @@ func (v *glmSearchConfigV0) ToGLMSearchConfig() (GLMSearchConfig, *GLMSearchSecu }, sec } +type baiduSearchConfigV0 struct { + Enabled bool `json:"enabled" env:"PICOCLAW_TOOLS_WEB_BAIDU_ENABLED"` + APIKey string `json:"api_key" env:"PICOCLAW_TOOLS_WEB_BAIDU_API_KEY"` + BaseURL string `json:"base_url" env:"PICOCLAW_TOOLS_WEB_BAIDU_BASE_URL"` + MaxResults int `json:"max_results" env:"PICOCLAW_TOOLS_WEB_BAIDU_MAX_RESULTS"` +} + +func (v *baiduSearchConfigV0) ToBaiduSearchConfig() (BaiduSearchConfig, *BaiduSearchSecurity) { + var sec *BaiduSearchSecurity + if v.APIKey != "" { + sec = &BaiduSearchSecurity{ + APIKey: v.APIKey, + } + } + return BaiduSearchConfig{ + Enabled: v.Enabled, + apiKey: v.APIKey, + BaseURL: v.BaseURL, + MaxResults: v.MaxResults, + }, sec +} + func (v *webToolsConfigV0) ToWebToolsConfig() (WebToolsConfig, WebToolsSecurity) { brave, braveSecurity := v.Brave.ToBraveConfig() tavily, tavilySecurity := v.Tavily.ToTavilyConfig() perplexity, perplexitySecurity := v.Perplexity.ToPerplexityConfig() glmSearch, glmSearchSecurity := v.GLMSearch.ToGLMSearchConfig() + baiduSearch, baiduSearchSecurity := v.BaiduSearch.ToBaiduSearchConfig() return WebToolsConfig{ ToolConfig: v.ToolConfig, @@ -1028,16 +962,18 @@ func (v *webToolsConfigV0) ToWebToolsConfig() (WebToolsConfig, WebToolsSecurity) Perplexity: perplexity, SearXNG: v.SearXNG, GLMSearch: glmSearch, + BaiduSearch: baiduSearch, PreferNative: v.PreferNative, Proxy: v.Proxy, FetchLimitBytes: v.FetchLimitBytes, Format: v.Format, PrivateHostWhitelist: v.PrivateHostWhitelist, }, WebToolsSecurity{ - Brave: braveSecurity, - Tavily: tavilySecurity, - Perplexity: perplexitySecurity, - GLMSearch: glmSearchSecurity, + Brave: braveSecurity, + Tavily: tavilySecurity, + Perplexity: perplexitySecurity, + GLMSearch: glmSearchSecurity, + BaiduSearch: baiduSearchSecurity, } } diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go index b356d474f..762988b0e 100644 --- a/pkg/config/config_test.go +++ b/pkg/config/config_test.go @@ -360,6 +360,96 @@ func TestSaveConfig_IncludesEmptyLegacyModelField(t *testing.T) { } } +func TestSaveConfig_PreservesDisabledTelegramPlaceholder(t *testing.T) { + tmpDir := t.TempDir() + path := filepath.Join(tmpDir, "config.json") + + cfg := DefaultConfig() + cfg.Channels.Telegram.Placeholder.Enabled = false + + if err := SaveConfig(path, cfg); err != nil { + t.Fatalf("SaveConfig failed: %v", err) + } + + data, err := os.ReadFile(path) + if err != nil { + t.Fatalf("ReadFile failed: %v", err) + } + if !strings.Contains(string(data), `"placeholder": {`) { + t.Fatalf("saved config should include telegram placeholder config, got: %s", string(data)) + } + if !strings.Contains(string(data), `"enabled": false`) { + t.Fatalf("saved config should persist placeholder.enabled=false, got: %s", string(data)) + } + + loaded, err := LoadConfig(path) + if err != nil { + t.Fatalf("LoadConfig failed: %v", err) + } + if loaded.Channels.Telegram.Placeholder.Enabled { + t.Fatal("telegram placeholder should remain disabled after SaveConfig/LoadConfig round-trip") + } +} + +// TestSaveConfig_FiltersVirtualModels verifies that SaveConfig does not write +// virtual models (generated by expandMultiKeyModels) to the config file. +func TestSaveConfig_FiltersVirtualModels(t *testing.T) { + tmpDir := t.TempDir() + path := filepath.Join(tmpDir, "config.json") + + cfg := DefaultConfig() + + // Manually add a virtual model to ModelList (simulating what expandMultiKeyModels does) + primaryModel := &ModelConfig{ + ModelName: "gpt-4", + Model: "openai/gpt-4o", + apiKeys: []string{"key1"}, + } + virtualModel := &ModelConfig{ + ModelName: "gpt-4__key_1", + Model: "openai/gpt-4o", + apiKeys: []string{"key2"}, + isVirtual: true, + } + cfg.ModelList = []*ModelConfig{primaryModel, virtualModel} + + // SaveConfig should filter out virtual models + if err := SaveConfig(path, cfg); err != nil { + t.Fatalf("SaveConfig failed: %v", err) + } + + // Reload and verify + reloaded, err := LoadConfig(path) + if err != nil { + t.Fatalf("LoadConfig failed: %v", err) + } + + // Should only have the primary model, not the virtual one + if len(reloaded.ModelList) != 1 { + t.Fatalf("expected 1 model after reload, got %d", len(reloaded.ModelList)) + } + + if reloaded.ModelList[0].ModelName != "gpt-4" { + t.Errorf("expected model_name 'gpt-4', got %q", reloaded.ModelList[0].ModelName) + } + + // Verify virtual model was not persisted + for _, m := range reloaded.ModelList { + if m.ModelName == "gpt-4__key_1" { + t.Errorf("virtual model gpt-4__key_1 should not have been saved") + } + } + + // Verify the saved file does not contain the virtual model name + data, err := os.ReadFile(path) + if err != nil { + t.Fatalf("ReadFile failed: %v", err) + } + if strings.Contains(string(data), "gpt-4__key_1") { + t.Errorf("saved config should not contain virtual model name 'gpt-4__key_1'") + } +} + // TestConfig_Complete verifies all config fields are set func TestConfig_Complete(t *testing.T) { cfg := DefaultConfig() @@ -397,6 +487,33 @@ func TestDefaultConfig_WebPreferNativeEnabled(t *testing.T) { } } +func TestDefaultConfig_ToolFeedbackDisabled(t *testing.T) { + cfg := DefaultConfig() + if cfg.Agents.Defaults.ToolFeedback.Enabled { + t.Fatal("DefaultConfig().Agents.Defaults.ToolFeedback.Enabled should be false") + } +} + +func TestLoadConfig_ToolFeedbackDefaultsFalseWhenUnset(t *testing.T) { + dir := t.TempDir() + configPath := filepath.Join(dir, "config.json") + if err := os.WriteFile( + configPath, + []byte(`{"version":1,"agents":{"defaults":{"workspace":"./workspace"}}}`), + 0o600, + ); err != nil { + t.Fatalf("WriteFile() error: %v", err) + } + + cfg, err := LoadConfig(configPath) + if err != nil { + t.Fatalf("LoadConfig() error: %v", err) + } + if cfg.Agents.Defaults.ToolFeedback.Enabled { + t.Fatal("agents.defaults.tool_feedback.enabled should remain false when unset in config file") + } +} + func TestLoadConfig_WebPreferNativeDefaultsTrueWhenUnset(t *testing.T) { dir := t.TempDir() configPath := filepath.Join(dir, "config.json") @@ -809,6 +926,91 @@ func TestFlexibleStringSlice_UnmarshalText_EmptySliceConsistency(t *testing.T) { }) } +func TestFlexibleStringSlice_UnmarshalJSON(t *testing.T) { + tests := []struct { + name string + input string + expected []string + }{ + { + name: "single string", + input: `"Thinking..."`, + expected: []string{"Thinking..."}, + }, + { + name: "single number", + input: `123`, + expected: []string{"123"}, + }, + { + name: "string array", + input: `["Thinking...", "Still working..."]`, + expected: []string{"Thinking...", "Still working..."}, + }, + { + name: "mixed array", + input: `["123", 456]`, + expected: []string{"123", "456"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var f FlexibleStringSlice + if err := json.Unmarshal([]byte(tt.input), &f); err != nil { + t.Fatalf("json.Unmarshal(%s) error = %v", tt.input, err) + } + if len(f) != len(tt.expected) { + t.Fatalf("json.Unmarshal(%s) len = %d, want %d", tt.input, len(f), len(tt.expected)) + } + for i, want := range tt.expected { + if f[i] != want { + t.Fatalf("json.Unmarshal(%s)[%d] = %q, want %q", tt.input, i, f[i], want) + } + } + }) + } +} + +func TestLoadConfig_TelegramPlaceholderTextAcceptsSingleString(t *testing.T) { + dir := t.TempDir() + cfgPath := filepath.Join(dir, "config.json") + data := `{ + "version": 1, + "agents": { "defaults": { "workspace": "", "model": "", "max_tokens": 0, "max_tool_iterations": 0 } }, + "bindings": [], + "session": {}, + "channels": { + "telegram": { + "enabled": true, + "bot_token": "", + "allow_from": [], + "placeholder": { + "enabled": true, + "text": "Thinking..." + } + } + }, + "model_list": [], + "gateway": {}, + "tools": {}, + "heartbeat": {}, + "devices": {}, + "voice": {} + }` + if err := os.WriteFile(cfgPath, []byte(data), 0o600); err != nil { + t.Fatalf("setup: %v", err) + } + + cfg, err := LoadConfig(cfgPath) + if err != nil { + t.Fatalf("LoadConfig() error = %v", err) + } + if got := []string(cfg.Channels.Telegram.Placeholder.Text); len(got) != 1 || got[0] != "Thinking..." { + t.Fatalf("placeholder.text = %#v, want [\"Thinking...\"]", got) + } +} + // TestLoadConfig_WarnsForPlaintextAPIKey verifies that LoadConfig resolves a plaintext // api_key into memory but does NOT rewrite the config file. File writes are the sole // responsibility of SaveConfig. @@ -1372,8 +1574,7 @@ func TestFilterSensitiveData_AllTokenTypes(t *testing.T) { Feishu: &FeishuSecurity{AppSecret: "feishu-app-secret-123", EncryptKey: "feishu-encrypt-key"}, DingTalk: &DingTalkSecurity{ClientSecret: "dingtalk-client-secret"}, OneBot: &OneBotSecurity{AccessToken: "onebot-access-token"}, - WeCom: &WeComSecurity{Token: "wecom-token", EncodingAESKey: "wecom-aes-key"}, - WeComApp: &WeComAppSecurity{CorpSecret: "wecom-app-secret", Token: "wecom-app-token"}, + WeCom: &WeComSecurity{Secret: "wecom-secret"}, Pico: &PicoSecurity{Token: "pico-token-abc123"}, IRC: &IRCSecurity{ Password: "irc-password", diff --git a/pkg/config/defaults.go b/pkg/config/defaults.go index c1d0ea0f6..b2148bbde 100644 --- a/pkg/config/defaults.go +++ b/pkg/config/defaults.go @@ -39,9 +39,10 @@ func DefaultConfig() *Config { SummarizeTokenPercent: 75, SteeringMode: "one-at-a-time", ToolFeedback: ToolFeedbackConfig{ - Enabled: true, + Enabled: false, MaxArgsLength: 300, }, + SplitOnMarker: false, }, }, Bindings: []AgentBinding{}, @@ -62,7 +63,7 @@ func DefaultConfig() *Config { Typing: TypingConfig{Enabled: true}, Placeholder: PlaceholderConfig{ Enabled: true, - Text: "Thinking... 💭", + Text: FlexibleStringSlice{"Thinking... 💭"}, }, Streaming: StreamingConfig{Enabled: true, ThrottleSeconds: 3, MinGrowthChars: 200}, UseMarkdownV2: false, @@ -111,8 +112,10 @@ func DefaultConfig() *Config { }, Placeholder: PlaceholderConfig{ Enabled: true, - Text: "Thinking... 💭", + Text: FlexibleStringSlice{"Thinking... 💭"}, }, + CryptoDatabasePath: "", + CryptoPassphrase: "", }, LINE: LINEConfig{ Enabled: false, @@ -129,32 +132,11 @@ func DefaultConfig() *Config { AllowFrom: FlexibleStringSlice{}, }, WeCom: WeComConfig{ - Enabled: false, - WebhookURL: "", - WebhookHost: "0.0.0.0", - WebhookPort: 18793, - WebhookPath: "/webhook/wecom", - AllowFrom: FlexibleStringSlice{}, - ReplyTimeout: 5, - }, - WeComApp: WeComAppConfig{ - Enabled: false, - CorpID: "", - AgentID: 0, - WebhookHost: "0.0.0.0", - WebhookPort: 18792, - WebhookPath: "/webhook/wecom-app", - AllowFrom: FlexibleStringSlice{}, - ReplyTimeout: 5, - }, - WeComAIBot: WeComAIBotConfig{ - Enabled: false, - WebhookPath: "/webhook/wecom-aibot", - AllowFrom: FlexibleStringSlice{}, - ReplyTimeout: 5, - MaxSteps: 10, - WelcomeMessage: "Hello! I'm your AI assistant. How can I help you today?", - ProcessingMessage: DefaultWeComAIBotProcessingMessage, + Enabled: false, + BotID: "", + WebSocketURL: "wss://openws.work.weixin.qq.com", + SendThinkingMessage: true, + AllowFrom: FlexibleStringSlice{}, }, Weixin: WeixinConfig{ Enabled: false, diff --git a/pkg/config/example_security_usage.go b/pkg/config/example_security_usage.go index cba76c6bc..42a1831b0 100644 --- a/pkg/config/example_security_usage.go +++ b/pkg/config/example_security_usage.go @@ -11,20 +11,33 @@ Package config # Example: Using Security Configuration -## 1. Create security.yml +## Overview -File: ~/.picoclaw/security.yml +The security configuration feature allows you to separate sensitive data (API keys, +tokens, secrets, passwords) from your main configuration. The system automatically +loads values from `.security.yml` and applies them to the corresponding fields in +your config. + +**Key Points:** +- Values from `.security.yml` are automatically mapped to config fields +- No `ref:` syntax is needed - just omit sensitive fields from config.json +- If a field exists in both files, `.security.yml` value takes precedence +- You can mix direct values in config.json with security values + +## 1. Create .security.yml + +File: ~/.picoclaw/.security.yml ```yaml # Model API Keys -# Note: Use 'api_keys' array for multiple keys (load balancing/failover) -# Single key should be provided as an array with one element +# All models MUST use 'api_keys' (plural) array format +# Even a single key must be provided as an array with one element model_list: gpt-5.4: api_keys: - "sk-proj-your-actual-openai-key-1" - - "sk-proj-your-actual-openai-key-2" # Failover key + - "sk-proj-your-actual-openai-key-2" # Optional: Multiple keys for failover claude-sonnet-4.6: api_keys: - "sk-ant-your-actual-anthropic-key" # Single key in array format @@ -38,80 +51,95 @@ channels: token: "your-discord-bot-token" # Web Tool Keys -# Note: Use 'api_keys' array for multiple keys (load balancing/failover) -# For GLMSearch, use 'api_key' (single string) +# Brave, Tavily, Perplexity: Use 'api_keys' array +# GLMSearch, BaiduSearch: Use 'api_key' single string web: brave: api_keys: - "BSAyour-brave-api-key-1" - - "BSAyour-brave-api-key-2" # Failover key + - "BSAyour-brave-api-key-2" # Optional: Multiple keys for failover tavily: api_keys: - "tvly-your-tavily-api-key" # Single key in array format + perplexity: + api_keys: + - "pplx-your-perplexity-api-key" # Single key in array format glm_search: api_key: "your-glm-search-api-key" # Single key (not array) + baidu_search: + api_key: "your-baidu-search-api-key" # Single key (not array) ``` -## 2. Update config.json to use references +## 2. Simplify config.json File: ~/.picoclaw/config.json +Note: Sensitive fields are omitted because they're loaded from .security.yml + ```json - { - "version": 1, - "agents": { - "defaults": { - "workspace": "~/picoclaw-workspace", - "model_name": "gpt-5.4" - } - }, - "model_list": [ - { - "model_name": "gpt-5.4", - "model": "openai/gpt-5.4", - "api_base": "https://api.openai.com/v1", - "api_key": "ref:model_list.gpt-5.4.api_key" - }, - { - "model_name": "claude-sonnet-4.6", - "model": "anthropic/claude-sonnet-4.6", - "api_base": "https://api.anthropic.com/v1", - "api_key": "ref:model_list.claude-sonnet-4.6.api_key" - } - ], - "channels": { - "telegram": { - "enabled": true, - "token": "ref:channels.telegram.token" - }, - "discord": { - "enabled": true, - "token": "ref:channels.discord.token" - } - }, + { + "version": 1, + "agents": { + "defaults": { + "workspace": "~/picoclaw-workspace", + "model_name": "gpt-5.4" + } + }, + "model_list": [ + { + "model_name": "gpt-5.4", + "model": "openai/gpt-5.4", + "api_base": "https://api.openai.com/v1" + // api_key is automatically loaded from .security.yml + }, + { + "model_name": "claude-sonnet-4.6", + "model": "anthropic/claude-sonnet-4.6", + "api_base": "https://api.anthropic.com/v1" + // api_key is automatically loaded from .security.yml + } + ], + "channels": { + "telegram": { + "enabled": true + // token is automatically loaded from .security.yml + }, + "discord": { + "enabled": true + // token is automatically loaded from .security.yml + } + }, "tools": { "web": { "brave": { - "enabled": true, - "api_key": "ref:web.brave.api_key" + "enabled": true + // api_key is automatically loaded from .security.yml }, "tavily": { - "enabled": true, - "api_key": "ref:web.tavily.api_key" + "enabled": true + // api_key is automatically loaded from .security.yml + }, + "glm_search": { + "enabled": true + // api_key is automatically loaded from .security.yml + }, + "baidu_search": { + "enabled": true + // api_key is automatically loaded from .security.yml } } } - } + } ``` ## 3. Set proper permissions ```bash -chmod 600 ~/.picoclaw/security.yml +chmod 600 ~/.picoclaw/.security.yml ``` ## 4. Add to .gitignore @@ -127,57 +155,131 @@ chmod 600 ~/.picoclaw/security.yml picoclaw --version ``` -# Available Reference Paths +# Supported Fields in .security.yml ## Model API Keys -- ref:model_list..api_key + +All models MUST use the `api_keys` (plural) array format in .security.yml. + +```yaml +model_list: + + : + api_keys: + - "key-1" + - "key-2" # Optional: Multiple keys for failover + +``` Examples: -- ref:model_list.gpt-5.4.api_key -- ref:model_list.claude-sonnet-4.6.api_key +```yaml +model_list: -**Note:** In .security.yml, use `api_keys` (array) format for models. -Both single and multiple keys should use the array format. + gpt-5.4: + api_keys: + - "sk-proj-key-1" + - "sk-proj-key-2" + claude-sonnet-4.6: + api_keys: + - "sk-ant-key" + +``` + +**Important:** +- Always use `api_keys` (plural) for models +- Even a single key must be in an array format +- The model_name in .security.yml must match the model_name in config.json ## Channel Tokens/Secrets -- ref:channels.telegram.token -- ref:channels.feishu.app_secret -- ref:channels.feishu.encrypt_key -- ref:channels.feishu.verification_token -- ref:channels.discord.token -- ref:channels.qq.app_secret -- ref:channels.dingtalk.client_secret -- ref:channels.slack.bot_token -- ref:channels.slack.app_token -- ref:channels.matrix.access_token -- ref:channels.line.channel_secret -- ref:channels.line.channel_access_token -- ref:channels.onebot.access_token -- ref:channels.wecom.token -- ref:channels.wecom.encoding_aes_key -- ref:channels.wecom_app.corp_secret -- ref:channels.wecom_app.token -- ref:channels.wecom_app.encoding_aes_key -- ref:channels.wecom_aibot.token -- ref:channels.wecom_aibot.encoding_aes_key -- ref:channels.pico.token -- ref:channels.irc.password -- ref:channels.irc.nickserv_password -- ref:channels.irc.sasl_password + +```yaml +channels: + + telegram: + token: "value" + feishu: + app_secret: "value" + encrypt_key: "value" + verification_token: "value" + discord: + token: "value" + weixin: + token: "value" + qq: + app_secret: "value" + dingtalk: + client_secret: "value" + slack: + bot_token: "value" + app_token: "value" + matrix: + access_token: "value" + line: + channel_secret: "value" + channel_access_token: "value" + onebot: + access_token: "value" + wecom: + token: "value" + encoding_aes_key: "value" + wecom_app: + corp_secret: "value" + token: "value" + encoding_aes_key: "value" + wecom_aibot: + secret: "value" + token: "value" + encoding_aes_key: "value" + pico: + token: "value" + irc: + password: "value" + nickserv_password: "value" + sasl_password: "value" ## Web Tool API Keys -- ref:web.brave.api_key -- ref:web.tavily.api_key -- ref:web.perplexity.api_key -- ref:web.glm_search.api_key -**Note:** -- Brave, Tavily, Perplexity: Use `api_keys` (array) format in .security.yml -- GLMSearch: Use `api_key` (single string) format in .security.yml +**Brave, Tavily, Perplexity:** +```yaml +web: + + brave: + api_keys: + - "BSA-key-1" + - "BSA-key-2" + tavily: + api_keys: + - "tvly-key" + perplexity: + api_keys: + - "pplx-key" + +``` +Use `api_keys` (plural) array format. + +**GLMSearch, BaiduSearch:** +```yaml +web: + + glm_search: + api_key: "your-glm-key" + baidu_search: + api_key: "your-baidu-key" + +``` +Use `api_key` (singular) single string format. ## Skills Registry Tokens -- ref:skills.github.token -- ref:skills.clawhub.auth_token + +```yaml +skills: + + github: + token: "value" + clawhub: + auth_token: "value" + +``` # Backward Compatibility @@ -191,14 +293,14 @@ You can still use direct values in config.json if needed: "model_name": "local-model", "model": "ollama/llama3", "api_base": "http://localhost:11434/v1", - "api_key": "ollama" // Direct value (no reference) + "api_key": "ollama" // Direct value (works fine) } ] } ``` -You can also mix references and direct values: +You can also mix security values and direct values: ```json @@ -206,10 +308,12 @@ You can also mix references and direct values: "model_list": [ { "model_name": "cloud-model", - "api_key": "ref:model_list.cloud-model.api_key" // From .security.yml + // api_key loaded from .security.yml }, { "model_name": "local-model", + "model": "ollama/llama3", + "api_base": "http://localhost:11434/v1", "api_key": "ollama" // Direct value } ] @@ -217,6 +321,11 @@ You can also mix references and direct values: ``` +**Priority Order:** +1. Environment variables (highest priority) +2. .security.yml values +3. config.json direct values (lowest priority) + # Migration from Old Config ## Step 1: Backup your config @@ -224,7 +333,7 @@ You can also mix references and direct values: cp ~/.picoclaw/config.json ~/.picoclaw/config.json.backup ``` -## Step 2: Copy the example security file +## Step 2: Create .security.yml ```bash cp security.example.yml ~/.picoclaw/.security.yml ``` @@ -232,10 +341,19 @@ cp security.example.yml ~/.picoclaw/.security.yml ## Step 3: Fill in your API keys Edit ~/.picoclaw/.security.yml and replace placeholders with your actual keys. -## Step 4: Update config.json references -Replace sensitive values in ~/.picoclaw/config.json with ref: references. +## Step 4: Simplify config.json (Recommended) +Remove sensitive fields from ~/.picoclaw/config.json: +- `api_key` fields from model_list entries +- `token` fields from channels +- `api_key` fields from tools.web +- `token`/`auth_token` fields from tools.skills -## Step 5: Test +## Step 5: Set permissions +```bash +chmod 600 ~/.picoclaw/.security.yml +``` + +## Step 6: Test ```bash picoclaw --version ``` @@ -249,9 +367,11 @@ rm ~/.picoclaw/config.json.backup ## Multiple API Keys (Load Balancing & Failover) -You can configure multiple API keys for both models and web tools to enable: +You can configure multiple API keys for models and web tools to enable: - **Load balancing**: Requests are distributed across multiple keys - **Failover**: If a key fails, the system automatically switches to another key +- **Rate limit management**: Distribute usage across multiple keys +- **High availability**: Reduce downtime during API provider issues ### Example: Model with Multiple Keys @@ -275,7 +395,7 @@ model_list: { "model_name": "gpt-5.4", "model": "openai/gpt-5.4", - "api_key": "ref:model_list.gpt-5.4.api_key" + "api_base": "https://api.openai.com/v1" } ] } @@ -307,8 +427,13 @@ web: "tools": { "web": { "brave": { - "enabled": true, - "api_key": "ref:web.brave.api_key" + "enabled": true + }, + "tavily": { + "enabled": true + }, + "glm_search": { + "enabled": true } } } @@ -316,9 +441,9 @@ web: ``` -### Single Key +## Single Key Format -Use array format with one element: +**Models, Brave, Tavily, Perplexity:** ```yaml model_list: @@ -328,36 +453,32 @@ model_list: ``` -### Multiple Keys (Load Balancing & Failover) - -Use array format with multiple elements: +**GLMSearch, BaiduSearch:** ```yaml -model_list: +web: - gpt-5.4: - api_keys: - - "sk-proj-key-1" - - "sk-proj-key-2" - - "sk-proj-key-3" + glm_search: + api_key: "your-glm-key" # Single key (not array) ``` -**Important:** All model keys in .security.yml must use the `api_keys` (plural) array format. -The single `api_key` (singular) format is NOT supported for models. - -### Model Index Matching +## Model Name Matching The system supports intelligent model name matching in .security.yml: -**Example 1: Exact Match** -```yaml -# config.json +### Example 1: Exact Match + +**config.json:** +```json { "model_name": "gpt-5.4:0" } -# .security.yml (exact match with index) +``` + +**.security.yml (exact match with index):** +```yaml model_list: gpt-5.4:0: @@ -365,26 +486,30 @@ model_list: ``` -**Example 2: Base Name Match** -```yaml -# config.json +### Example 2: Base Name Match + +**config.json:** +```json { "model_name": "gpt-5.4:0" } -# .security.yml (base name without index) +``` + +**.security.yml (base name without index):** +```yaml model_list: gpt-5.4: - api_keys: ["key-1"] + api_keys: ["key-1", "key-2"] ``` Both methods work. The base name match allows you to use simpler keys in .security.yml even when your config uses indexed model names for load balancing. -### Security File Permissions +## Security File Permissions The security file should have restricted permissions: @@ -397,26 +522,64 @@ This ensures only the owner can read and write the file. # Security Best Practices 1. Never commit .security.yml to version control -2. Set file permissions: chmod 600 ~/.picoclaw/.security.yml -3. Use different keys for different environments -4. Rotate keys regularly and update .security.yml -5. Encrypt backups containing .security.yml +2. Add .security.yml to your .gitignore file +3. Set file permissions: chmod 600 ~/.picoclaw/.security.yml +4. Use different keys for different environments (dev, staging, production) +5. Rotate keys regularly and update .security.yml +6. Encrypt backups containing .security.yml +7. Review access regularly + +# Environment Variables + +You can override any security value using environment variables: + +```bash +# Channels +export PICOCLAW_CHANNELS_TELEGRAM_TOKEN="token-from-env" +export PICOCLAW_CHANNELS_DISCORD_TOKEN="discord-token-from-env" + +# Web Tools +export PICOCLAW_TOOLS_WEB_BRAVE_API_KEY="brave-key-from-env" +export PICOCLAW_TOOLS_WEB_BAIDU_API_KEY="baidu-key-from-env" + +# Skills +export PICOCLAW_TOOLS_SKILLS_GITHUB_TOKEN="github-token-from-env" +``` + +Environment variables have the highest priority and will override both config.json +and .security.yml values. # Troubleshooting +## Error: "failed to load security config" +- Ensure .security.yml exists in the same directory as config.json +- Check YAML syntax is valid (use a YAML validator) +- Verify file permissions allow reading + ## Error: "model security entry not found" - Check that the model name in config.json matches exactly in .security.yml - Verify the model_list section exists in .security.yml +- For indexed names (e.g., "gpt-5.4:0"), check both exact match and base name match +- Ensure the YAML structure is correct (proper indentation) -## Error: "failed to load security config" -- Ensure .security.yml exists in the same directory as config.json -- Check YAML syntax is valid -- Verify file permissions allow reading +## Multiple API Keys Not Working +- Ensure you're using `api_keys` (plural) in .security.yml for models and web tools (except GLMSearch/BaiduSearch) +- Check that the array format is correct in YAML (proper indentation with dashes) +- Remember: Models, Brave, Tavily, Perplexity MUST use `api_keys` (array format) +- GLMSearch and BaiduSearch MUST use `api_key` (single string format) -## Error: "unknown reference path" -- Verify the reference format is correct -- Check the path structure matches the examples above -- Ensure all required sections exist in .security.yml +## Keys Not Being Applied +- Check that .security.yml is in the same directory as config.json +- Verify the file permissions allow reading (chmod 600 ~/.picoclaw/.security.yml) +- Ensure the YAML structure matches the expected format +- Check for typos in field names (case-sensitive) +- Verify the model/channel names match exactly (case-sensitive) + +## Load Balancing/Failover Issues +- Verify all API keys in the api_keys array are valid +- Check that all keys have the same rate limits and permissions +- Monitor logs to see which keys are being used and failing +- Ensure the api_keys array is properly formatted in YAML */ package config diff --git a/pkg/config/multikey_test.go b/pkg/config/multikey_test.go index cc529905c..c17fcc53b 100644 --- a/pkg/config/multikey_test.go +++ b/pkg/config/multikey_test.go @@ -232,6 +232,78 @@ func TestExpandMultiKeyModels_PreservesOtherFields(t *testing.T) { } } +func TestExpandMultiKeyModels_IsVirtualFlag(t *testing.T) { + models := []*ModelConfig{ + { + ModelName: "gpt-4", + Model: "openai/gpt-4o", + apiKeys: []string{"key1", "key2", "key3"}, + }, + } + + result := expandMultiKeyModels(models) + + // Should expand to 3 models + if len(result) != 3 { + t.Fatalf("expected 3 models, got %d", len(result)) + } + + // Primary model should NOT be virtual + primary := result[2] + if primary.isVirtual { + t.Errorf("primary model should not be virtual") + } + if primary.ModelName != "gpt-4" { + t.Errorf("expected primary model_name 'gpt-4', got %q", primary.ModelName) + } + + // Virtual models should have isVirtual = true + virtual1 := result[0] + if !virtual1.isVirtual { + t.Errorf("gpt-4__key_1 should be virtual") + } + if virtual1.ModelName != "gpt-4__key_1" { + t.Errorf("expected virtual model_name 'gpt-4__key_1', got %q", virtual1.ModelName) + } + + virtual2 := result[1] + if !virtual2.isVirtual { + t.Errorf("gpt-4__key_2 should be virtual") + } + if virtual2.ModelName != "gpt-4__key_2" { + t.Errorf("expected virtual model_name 'gpt-4__key_2', got %q", virtual2.ModelName) + } + + // IsVirtual() method should work + if !virtual1.IsVirtual() { + t.Errorf("IsVirtual() should return true for virtual model") + } + if primary.IsVirtual() { + t.Errorf("IsVirtual() should return false for primary model") + } +} + +func TestExpandMultiKeyModels_SingleKey_NotVirtual(t *testing.T) { + models := []*ModelConfig{ + { + ModelName: "gpt-4", + Model: "openai/gpt-4o", + apiKeys: []string{"single-key"}, + }, + } + + result := expandMultiKeyModels(models) + + if len(result) != 1 { + t.Fatalf("expected 1 model, got %d", len(result)) + } + + // Single key model should NOT be virtual + if result[0].isVirtual { + t.Errorf("single key model should not be virtual") + } +} + func TestMergeAPIKeys(t *testing.T) { tests := []struct { name string diff --git a/pkg/config/security.go b/pkg/config/security.go index da989ca88..47ad1a5b0 100644 --- a/pkg/config/security.go +++ b/pkg/config/security.go @@ -69,21 +69,19 @@ type ModelSecurityEntry struct { // ChannelsSecurity stores channel-related security data type ChannelsSecurity struct { - Telegram *TelegramSecurity `yaml:"telegram,omitempty"` - Feishu *FeishuSecurity `yaml:"feishu,omitempty"` - Discord *DiscordSecurity `yaml:"discord,omitempty"` - Weixin *WeixinSecurity `yaml:"weixin,omitempty"` - QQ *QQSecurity `yaml:"qq,omitempty"` - DingTalk *DingTalkSecurity `yaml:"dingtalk,omitempty"` - Slack *SlackSecurity `yaml:"slack,omitempty"` - Matrix *MatrixSecurity `yaml:"matrix,omitempty"` - LINE *LINESecurity `yaml:"line,omitempty"` - OneBot *OneBotSecurity `yaml:"onebot,omitempty"` - WeCom *WeComSecurity `yaml:"wecom,omitempty"` - WeComApp *WeComAppSecurity `yaml:"wecom_app,omitempty"` - WeComAIBot *WeComAIBotSecurity `yaml:"wecom_aibot,omitempty"` - Pico *PicoSecurity `yaml:"pico,omitempty"` - IRC *IRCSecurity `yaml:"irc,omitempty"` + Telegram *TelegramSecurity `yaml:"telegram,omitempty"` + Feishu *FeishuSecurity `yaml:"feishu,omitempty"` + Discord *DiscordSecurity `yaml:"discord,omitempty"` + Weixin *WeixinSecurity `yaml:"weixin,omitempty"` + QQ *QQSecurity `yaml:"qq,omitempty"` + DingTalk *DingTalkSecurity `yaml:"dingtalk,omitempty"` + Slack *SlackSecurity `yaml:"slack,omitempty"` + Matrix *MatrixSecurity `yaml:"matrix,omitempty"` + LINE *LINESecurity `yaml:"line,omitempty"` + OneBot *OneBotSecurity `yaml:"onebot,omitempty"` + WeCom *WeComSecurity `yaml:"wecom,omitempty"` + Pico *PicoSecurity `yaml:"pico,omitempty"` + IRC *IRCSecurity `yaml:"irc,omitempty"` } type TelegramSecurity struct { @@ -131,20 +129,7 @@ type OneBotSecurity struct { } type WeComSecurity struct { - Token string `yaml:"token,omitempty" env:"PICOCLAW_CHANNELS_WECOM_TOKEN"` - EncodingAESKey string `yaml:"encoding_aes_key,omitempty" env:"PICOCLAW_CHANNELS_WECOM_ENCODING_AES_KEY"` -} - -type WeComAppSecurity struct { - CorpSecret string `yaml:"corp_secret,omitempty" env:"PICOCLAW_CHANNELS_WECOM_APP_CORP_SECRET"` - Token string `yaml:"token,omitempty" env:"PICOCLAW_CHANNELS_WECOM_APP_TOKEN"` - EncodingAESKey string `yaml:"encoding_aes_key,omitempty" env:"PICOCLAW_CHANNELS_WECOM_APP_ENCODING_AES_KEY"` -} - -type WeComAIBotSecurity struct { - Secret string `yaml:"secret,omitempty" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_SECRET"` - Token string `yaml:"token,omitempty" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_TOKEN"` - EncodingAESKey string `yaml:"encoding_aes_key,omitempty" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_ENCODING_AES_KEY"` + Secret string `yaml:"secret,omitempty" env:"PICOCLAW_CHANNELS_WECOM_SECRET"` } type PicoSecurity struct { @@ -334,17 +319,9 @@ func mergeChannelsSecurity(dst, src *ChannelsSecurity) { if src.OneBot != nil && src.OneBot.AccessToken != "" { dst.OneBot = src.OneBot } - if src.WeCom != nil && (src.WeCom.Token != "" || src.WeCom.EncodingAESKey != "") { + if src.WeCom != nil && src.WeCom.Secret != "" { dst.WeCom = src.WeCom } - if src.WeComApp != nil && - (src.WeComApp.CorpSecret != "" || src.WeComApp.Token != "" || src.WeComApp.EncodingAESKey != "") { - dst.WeComApp = src.WeComApp - } - if src.WeComAIBot != nil && - (src.WeComAIBot.Secret != "" || src.WeComAIBot.Token != "" || src.WeComAIBot.EncodingAESKey != "") { - dst.WeComAIBot = src.WeComAIBot - } if src.Pico != nil && src.Pico.Token != "" { dst.Pico = src.Pico } diff --git a/pkg/config/security_integration_test.go b/pkg/config/security_integration_test.go index 218914590..002988f2f 100644 --- a/pkg/config/security_integration_test.go +++ b/pkg/config/security_integration_test.go @@ -43,7 +43,8 @@ func TestSecurityConfigIntegration(t *testing.T) { t.Run("Full workflow with security references", func(t *testing.T) { tmpDir := t.TempDir() - // Create config.json with references + // Create config.json with direct security values (not ref: references) + // These values should take precedence over .security.yml configPath := filepath.Join(tmpDir, "config.json") configContent := `{ "version": 1, @@ -52,25 +53,25 @@ func TestSecurityConfigIntegration(t *testing.T) { "model_name": "test-model", "model": "openai/test-model", "api_base": "https://api.openai.com/v1", - "api_key": "ref:model_list.test-model.api_key" + "api_key": "sk-from-config-json-direct" } ], "channels": { "telegram": { "enabled": true, - "token": "ref:channels.telegram.token" + "token": "token-from-config-json-direct" } }, "tools": { "web": { "brave": { "enabled": true, - "api_key": "ref:web.brave.api_key" + "api_key": "BSA-from-config-json-direct" } }, "skills": { "github": { - "token": "ref:skills.github.token" + "token": "ghp-from-config-json-direct" } } } @@ -78,46 +79,47 @@ func TestSecurityConfigIntegration(t *testing.T) { err := os.WriteFile(configPath, []byte(configContent), 0o644) require.NoError(t, err) - // Create .security.yml with actual values + // Create .security.yml with different values + // These should be overridden by config.json values securityPath := filepath.Join(tmpDir, SecurityConfigFile) securityContent := `model_list: test-model: api_keys: - - "sk-test-api-key-12345" + - "sk-from-security-yml" channels: telegram: - token: "123456789:ABCdefGHIjklMNOpqrsTUVwxyz" + token: "token-from-security-yml" web: brave: api_keys: - - "BSAbrave-api-key-67890" + - "BSA-from-security-yml" skills: github: - token: "ghp_github-token-abc123"` + token: "ghp-from-security-yml"` err = os.WriteFile(securityPath, []byte(securityContent), 0o600) require.NoError(t, err) - // Load config and verify references are resolved + // Load config and verify config.json values take precedence cfg, err := LoadConfig(configPath) require.NoError(t, err) require.NotNil(t, cfg) - // Verify model API key is resolved + // Verify model API key from config.json takes precedence assert.Equal(t, 1, len(cfg.ModelList)) assert.Equal(t, "test-model", cfg.ModelList[0].ModelName) - assert.Equal(t, "sk-test-api-key-12345", cfg.ModelList[0].apiKeys[0]) + assert.Equal(t, "sk-from-config-json-direct", cfg.ModelList[0].apiKeys[0]) - // Verify channel token is resolved - assert.Equal(t, "123456789:ABCdefGHIjklMNOpqrsTUVwxyz", cfg.Channels.Telegram.token) + // Verify channel token from config.json takes precedence + assert.Equal(t, "token-from-config-json-direct", cfg.Channels.Telegram.token) - // Verify web tool API key is resolved - assert.Equal(t, "BSAbrave-api-key-67890", cfg.Tools.Web.Brave.APIKey()) + // Verify web tool API key from config.json takes precedence + assert.Equal(t, "BSA-from-config-json-direct", cfg.Tools.Web.Brave.APIKey()) - // Verify skills token is resolved - assert.Equal(t, "ghp_github-token-abc123", cfg.Tools.Skills.Github.token) + // Verify skills token from config.json takes precedence + assert.Equal(t, "ghp-from-config-json-direct", cfg.Tools.Skills.Github.token) }) } @@ -240,15 +242,7 @@ func TestAllSecurityKeysAccessible(t *testing.T) { }, "wecom": { "enabled": true, - "webhook_url": "https://qyapi.weixin.qq.com/cgi-bin/webhook" - }, - "wecom_app": { - "enabled": true, - "corp_id": "test_corp_id", - "agent_id": 123456 - }, - "wecom_aibot": { - "enabled": true + "bot_id": "test_wecom_bot_id" }, "pico": { "enabled": true @@ -315,15 +309,7 @@ channels: onebot: access_token: "onebot_test_access_token" wecom: - token: "wecom_test_webhook_token" - encoding_aes_key: "wecom_test_aes_key" - wecom_app: - corp_secret: "wecom_app_test_corp_secret" - token: "wecom_app_test_token" - encoding_aes_key: "wecom_app_test_aes_key" - wecom_aibot: - token: "wecom_aibot_test_token" - encoding_aes_key: "wecom_aibot_test_aes_key" + secret: "wecom_test_secret" pico: token: "pico_test_token" irc: @@ -409,24 +395,10 @@ skills: t.Logf("OneBot AccessToken(): %s", cfg.Channels.OneBot.AccessToken()) // WeCom - assert.Equal(t, "wecom_test_webhook_token", cfg.Channels.WeCom.Token()) - assert.Equal(t, "wecom_test_aes_key", cfg.Channels.WeCom.EncodingAESKey()) - t.Logf("WeCom Token(): %s", cfg.Channels.WeCom.Token()) - t.Logf("WeCom EncodingAESKey(): %s", cfg.Channels.WeCom.EncodingAESKey()) - - // WeCom App - assert.Equal(t, "wecom_app_test_corp_secret", cfg.Channels.WeComApp.CorpSecret()) - assert.Equal(t, "wecom_app_test_token", cfg.Channels.WeComApp.Token()) - assert.Equal(t, "wecom_app_test_aes_key", cfg.Channels.WeComApp.EncodingAESKey()) - t.Logf("WeComApp CorpSecret(): %s", cfg.Channels.WeComApp.CorpSecret()) - t.Logf("WeComApp Token(): %s", cfg.Channels.WeComApp.Token()) - t.Logf("WeComApp EncodingAESKey(): %s", cfg.Channels.WeComApp.EncodingAESKey()) - - // WeCom AI Bot - assert.Equal(t, "wecom_aibot_test_token", cfg.Channels.WeComAIBot.Token()) - assert.Equal(t, "wecom_aibot_test_aes_key", cfg.Channels.WeComAIBot.EncodingAESKey()) - t.Logf("WeComAIBot Token(): %s", cfg.Channels.WeComAIBot.Token()) - t.Logf("WeComAIBot EncodingAESKey(): %s", cfg.Channels.WeComAIBot.EncodingAESKey()) + assert.Equal(t, "test_wecom_bot_id", cfg.Channels.WeCom.BotID) + assert.Equal(t, "wecom_test_secret", cfg.Channels.WeCom.Secret()) + t.Logf("WeCom BotID: %s", cfg.Channels.WeCom.BotID) + t.Logf("WeCom Secret(): %s", cfg.Channels.WeCom.Secret()) // Pico assert.Equal(t, "pico_test_token", cfg.Channels.Pico.Token()) diff --git a/pkg/fileutil/file_test.go b/pkg/fileutil/file_test.go new file mode 100644 index 000000000..b0494d0d3 --- /dev/null +++ b/pkg/fileutil/file_test.go @@ -0,0 +1,176 @@ +package fileutil + +import ( + "os" + "path/filepath" + "sync" + "testing" +) + +func TestWriteFileAtomic_Basic(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "test.txt") + data := []byte("hello picoclaw") + + err := WriteFileAtomic(path, data, 0o644) + if err != nil { + t.Fatalf("WriteFileAtomic failed: %v", err) + } + + got, err := os.ReadFile(path) + if err != nil { + t.Fatalf("ReadFile failed: %v", err) + } + if string(got) != string(data) { + t.Errorf("got %q, want %q", got, data) + } +} + +func TestWriteFileAtomic_Permissions(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "secret.txt") + + err := WriteFileAtomic(path, []byte("secret"), 0o600) + if err != nil { + t.Fatalf("WriteFileAtomic failed: %v", err) + } + + info, err := os.Stat(path) + if err != nil { + t.Fatalf("Stat failed: %v", err) + } + // On Unix, check file mode (ignoring directory bits) + if got := info.Mode().Perm(); got != 0o600 { + t.Errorf("permissions = %o, want %o", got, 0o600) + } +} + +func TestWriteFileAtomic_Overwrite(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "overwrite.txt") + + // Write initial content + if err := WriteFileAtomic(path, []byte("old"), 0o644); err != nil { + t.Fatalf("first write failed: %v", err) + } + + // Overwrite + if err := WriteFileAtomic(path, []byte("new"), 0o644); err != nil { + t.Fatalf("second write failed: %v", err) + } + + got, _ := os.ReadFile(path) + if string(got) != "new" { + t.Errorf("got %q after overwrite, want %q", got, "new") + } +} + +func TestWriteFileAtomic_EmptyData(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "empty.txt") + + err := WriteFileAtomic(path, []byte{}, 0o644) + if err != nil { + t.Fatalf("WriteFileAtomic with empty data failed: %v", err) + } + + got, _ := os.ReadFile(path) + if len(got) != 0 { + t.Errorf("expected empty file, got %d bytes", len(got)) + } +} + +func TestWriteFileAtomic_CreatesParentDirs(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "a", "b", "c", "deep.txt") + + err := WriteFileAtomic(path, []byte("deep"), 0o644) + if err != nil { + t.Fatalf("WriteFileAtomic with nested dirs failed: %v", err) + } + + got, _ := os.ReadFile(path) + if string(got) != "deep" { + t.Errorf("got %q, want %q", got, "deep") + } +} + +func TestWriteFileAtomic_NoTempFileOnSuccess(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "clean.txt") + + if err := WriteFileAtomic(path, []byte("data"), 0o644); err != nil { + t.Fatalf("WriteFileAtomic failed: %v", err) + } + + // Verify no temp files remain + entries, _ := os.ReadDir(dir) + for _, e := range entries { + if e.Name() != "clean.txt" { + t.Errorf("unexpected file remaining: %s", e.Name()) + } + } +} + +func TestWriteFileAtomic_LargeFile(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "large.bin") + + // 1MB of data + data := make([]byte, 1<<20) + for i := range data { + data[i] = byte(i % 256) + } + + if err := WriteFileAtomic(path, data, 0o644); err != nil { + t.Fatalf("WriteFileAtomic with large file failed: %v", err) + } + + got, _ := os.ReadFile(path) + if len(got) != len(data) { + t.Errorf("file size = %d, want %d", len(got), len(data)) + } +} + +func TestWriteFileAtomic_Concurrent(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "concurrent.txt") + + var wg sync.WaitGroup + errs := make(chan error, 10) + + for i := 0; i < 10; i++ { + wg.Add(1) + go func(n int) { + defer wg.Done() + data := []byte(string(rune('A' + n))) + if err := WriteFileAtomic(path, data, 0o644); err != nil { + errs <- err + } + }(i) + } + + wg.Wait() + close(errs) + + for err := range errs { + t.Errorf("concurrent write error: %v", err) + } + + // File should exist and contain exactly 1 byte (last writer wins) + got, err := os.ReadFile(path) + if err != nil { + t.Fatalf("ReadFile after concurrent writes failed: %v", err) + } + if len(got) != 1 { + t.Errorf("expected 1 byte after concurrent writes, got %d", len(got)) + } +} + +func TestWriteFileAtomic_InvalidPath(t *testing.T) { + // /dev/null/impossible is not a valid path on any OS + err := WriteFileAtomic("/dev/null/impossible/file.txt", []byte("data"), 0o644) + if err == nil { + t.Error("expected error for invalid path, got nil") + } +} diff --git a/pkg/gateway/channel_matrix.go b/pkg/gateway/channel_matrix.go new file mode 100644 index 000000000..a46addae1 --- /dev/null +++ b/pkg/gateway/channel_matrix.go @@ -0,0 +1,24 @@ +//go:build !mipsle && !netbsd && !(freebsd && arm) + +package gateway + +import ( + // Matrix currently pulls in mautrix crypto and modernc sqlite transitively. + // + // We exclude it on: + // - linux/mipsle: mautrix crypto falls back to libolm when the `goolm` build + // tag is unavailable, and modernc.org/sqlite/modernc.org/libc also lacks a + // working build path for our mipsle + softfloat target. + // - netbsd/*: modernc.org/sqlite v1.46.1 fails to compile due to broken + // generated mutex code on NetBSD (for example sqlite_netbsd_amd64.go calls + // mu.enter/mu.leave, but the generated mutex type does not define them). + // - freebsd/arm: modernc.org/libc v1.67.6 fails to compile due to broken + // generated 32-bit FreeBSD code (size_t/uint64 and int32/int64 mismatches + // in libc_freebsd.go). + // + // This means Matrix is currently unavailable on those targets. The proper + // long-term fix is to split Matrix basic support from its E2EE/sqlite-backed + // crypto path, or to upgrade/replace the upstream sqlite dependency once the + // affected targets are supported. + _ "github.com/sipeed/picoclaw/pkg/channels/matrix" +) diff --git a/pkg/gateway/gateway.go b/pkg/gateway/gateway.go index fc2465747..03d7dfe0c 100644 --- a/pkg/gateway/gateway.go +++ b/pkg/gateway/gateway.go @@ -20,7 +20,6 @@ import ( _ "github.com/sipeed/picoclaw/pkg/channels/irc" _ "github.com/sipeed/picoclaw/pkg/channels/line" _ "github.com/sipeed/picoclaw/pkg/channels/maixcam" - _ "github.com/sipeed/picoclaw/pkg/channels/matrix" _ "github.com/sipeed/picoclaw/pkg/channels/onebot" _ "github.com/sipeed/picoclaw/pkg/channels/pico" _ "github.com/sipeed/picoclaw/pkg/channels/qq" diff --git a/pkg/health/server_test.go b/pkg/health/server_test.go new file mode 100644 index 000000000..6e0b5e66b --- /dev/null +++ b/pkg/health/server_test.go @@ -0,0 +1,347 @@ +package health + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "testing" + "time" +) + +func newTestServer() *Server { + s := &Server{ + ready: false, + checks: make(map[string]Check), + startTime: time.Now(), + } + return s +} + +func TestHealthHandler_ReturnsOK(t *testing.T) { + s := newTestServer() + req := httptest.NewRequest(http.MethodGet, "/health", nil) + w := httptest.NewRecorder() + + s.healthHandler(w, req) + + if w.Code != http.StatusOK { + t.Errorf("health status = %d, want %d", w.Code, http.StatusOK) + } + + var resp StatusResponse + if err := json.NewDecoder(w.Body).Decode(&resp); err != nil { + t.Fatalf("failed to decode response: %v", err) + } + if resp.Status != "ok" { + t.Errorf("status = %q, want %q", resp.Status, "ok") + } + if resp.Pid == 0 { + t.Error("pid should not be 0") + } + if resp.Uptime == "" { + t.Error("uptime should not be empty") + } +} + +func TestReadyHandler_NotReady(t *testing.T) { + s := newTestServer() + // s.ready defaults to false + req := httptest.NewRequest(http.MethodGet, "/ready", nil) + w := httptest.NewRecorder() + + s.readyHandler(w, req) + + if w.Code != http.StatusServiceUnavailable { + t.Errorf("ready status = %d, want %d", w.Code, http.StatusServiceUnavailable) + } + + var resp StatusResponse + if err := json.NewDecoder(w.Body).Decode(&resp); err != nil { + t.Fatalf("failed to decode response: %v", err) + } + if resp.Status != "not ready" { + t.Errorf("status = %q, want %q", resp.Status, "not ready") + } +} + +func TestReadyHandler_Ready(t *testing.T) { + s := newTestServer() + s.SetReady(true) + + req := httptest.NewRequest(http.MethodGet, "/ready", nil) + w := httptest.NewRecorder() + + s.readyHandler(w, req) + + if w.Code != http.StatusOK { + t.Errorf("ready status = %d, want %d", w.Code, http.StatusOK) + } + + var resp StatusResponse + if err := json.NewDecoder(w.Body).Decode(&resp); err != nil { + t.Fatalf("failed to decode response: %v", err) + } + if resp.Status != "ready" { + t.Errorf("status = %q, want %q", resp.Status, "ready") + } +} + +func TestReadyHandler_FailedCheck(t *testing.T) { + s := newTestServer() + s.SetReady(true) + + // Register a failing check + s.RegisterCheck("database", func() (bool, string) { + return false, "connection refused" + }) + + req := httptest.NewRequest(http.MethodGet, "/ready", nil) + w := httptest.NewRecorder() + + s.readyHandler(w, req) + + if w.Code != http.StatusServiceUnavailable { + t.Errorf("ready with failed check = %d, want %d", w.Code, http.StatusServiceUnavailable) + } + + var resp StatusResponse + if err := json.NewDecoder(w.Body).Decode(&resp); err != nil { + t.Fatalf("failed to decode response: %v", err) + } + if resp.Status != "not ready" { + t.Errorf("status = %q, want %q", resp.Status, "not ready") + } + check, ok := resp.Checks["database"] + if !ok { + t.Fatal("missing database check in response") + } + if check.Status != "fail" { + t.Errorf("check status = %q, want %q", check.Status, "fail") + } + if check.Message != "connection refused" { + t.Errorf("check message = %q, want %q", check.Message, "connection refused") + } +} + +func TestReadyHandler_PassingCheck(t *testing.T) { + s := newTestServer() + s.SetReady(true) + + s.RegisterCheck("redis", func() (bool, string) { + return true, "connected" + }) + + req := httptest.NewRequest(http.MethodGet, "/ready", nil) + w := httptest.NewRecorder() + + s.readyHandler(w, req) + + if w.Code != http.StatusOK { + t.Errorf("ready with passing check = %d, want %d", w.Code, http.StatusOK) + } + + var resp StatusResponse + if err := json.NewDecoder(w.Body).Decode(&resp); err != nil { + t.Fatalf("failed to decode response: %v", err) + } + if resp.Checks["redis"].Status != "ok" { + t.Errorf("redis check status = %q, want %q", resp.Checks["redis"].Status, "ok") + } +} + +func TestReloadHandler_MethodNotAllowed(t *testing.T) { + s := newTestServer() + + req := httptest.NewRequest(http.MethodGet, "/reload", nil) + w := httptest.NewRecorder() + + s.reloadHandler(w, req) + + if w.Code != http.StatusMethodNotAllowed { + t.Errorf("reload GET status = %d, want %d", w.Code, http.StatusMethodNotAllowed) + } +} + +func TestReloadHandler_NoReloadFunc(t *testing.T) { + s := newTestServer() + + req := httptest.NewRequest(http.MethodPost, "/reload", nil) + w := httptest.NewRecorder() + + s.reloadHandler(w, req) + + if w.Code != http.StatusServiceUnavailable { + t.Errorf("reload without func = %d, want %d", w.Code, http.StatusServiceUnavailable) + } +} + +func TestReloadHandler_Success(t *testing.T) { + s := newTestServer() + called := false + s.SetReloadFunc(func() error { + called = true + return nil + }) + + req := httptest.NewRequest(http.MethodPost, "/reload", nil) + w := httptest.NewRecorder() + + s.reloadHandler(w, req) + + if w.Code != http.StatusOK { + t.Errorf("reload status = %d, want %d", w.Code, http.StatusOK) + } + if !called { + t.Error("reload function was not called") + } +} + +func TestReloadHandler_Error(t *testing.T) { + s := newTestServer() + s.SetReloadFunc(func() error { + return errors.New("config parse error") + }) + + req := httptest.NewRequest(http.MethodPost, "/reload", nil) + w := httptest.NewRecorder() + + s.reloadHandler(w, req) + + if w.Code != http.StatusInternalServerError { + t.Errorf("reload error status = %d, want %d", w.Code, http.StatusInternalServerError) + } +} + +func TestSetReady_Toggle(t *testing.T) { + s := newTestServer() + + s.SetReady(true) + req := httptest.NewRequest(http.MethodGet, "/ready", nil) + w := httptest.NewRecorder() + s.readyHandler(w, req) + if w.Code != http.StatusOK { + t.Errorf("after SetReady(true): status = %d, want %d", w.Code, http.StatusOK) + } + + s.SetReady(false) + w = httptest.NewRecorder() + s.readyHandler(w, httptest.NewRequest(http.MethodGet, "/ready", nil)) + if w.Code != http.StatusServiceUnavailable { + t.Errorf("after SetReady(false): status = %d, want %d", w.Code, http.StatusServiceUnavailable) + } +} + +func TestRegisterCheck_MultipleChecks(t *testing.T) { + s := newTestServer() + s.SetReady(true) + + s.RegisterCheck("db", func() (bool, string) { + return true, "ok" + }) + s.RegisterCheck("cache", func() (bool, string) { + return true, "ok" + }) + s.RegisterCheck("queue", func() (bool, string) { + return false, "timeout" + }) + + req := httptest.NewRequest(http.MethodGet, "/ready", nil) + w := httptest.NewRecorder() + s.readyHandler(w, req) + + // Should be not ready because queue check fails + if w.Code != http.StatusServiceUnavailable { + t.Errorf("status = %d, want %d (queue check failed)", w.Code, http.StatusServiceUnavailable) + } + + var resp StatusResponse + if err := json.NewDecoder(w.Body).Decode(&resp); err != nil { + t.Fatalf("failed to decode response: %v", err) + } + if len(resp.Checks) != 3 { + t.Errorf("checks count = %d, want 3", len(resp.Checks)) + } +} + +func TestRegisterOnMux(t *testing.T) { + s := newTestServer() + s.SetReady(true) + + mux := http.NewServeMux() + s.RegisterOnMux(mux) + + // Test /health on custom mux + req := httptest.NewRequest(http.MethodGet, "/health", nil) + w := httptest.NewRecorder() + mux.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Errorf("/health on custom mux = %d, want %d", w.Code, http.StatusOK) + } + + // Test /ready on custom mux + req = httptest.NewRequest(http.MethodGet, "/ready", nil) + w = httptest.NewRecorder() + mux.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Errorf("/ready on custom mux = %d, want %d", w.Code, http.StatusOK) + } +} + +func TestNewServer(t *testing.T) { + s := NewServer("127.0.0.1", 0) + if s == nil { + t.Fatal("NewServer returned nil") + } + if s.ready { + t.Error("new server should not be ready by default") + } + if s.checks == nil { + t.Error("checks map should be initialized") + } +} + +func TestStartContext_Cancellation(t *testing.T) { + s := NewServer("127.0.0.1", 0) + + ctx, cancel := context.WithCancel(context.Background()) + + errCh := make(chan error, 1) + go func() { + errCh <- s.StartContext(ctx) + }() + + // Give server time to start + time.Sleep(50 * time.Millisecond) + + // Cancel context should trigger shutdown + cancel() + + select { + case err := <-errCh: + if err != nil { + t.Errorf("StartContext returned unexpected error: %v", err) + } + case <-time.After(2 * time.Second): + t.Error("StartContext did not return after context cancellation") + } +} + +func TestStatusString(t *testing.T) { + tests := []struct { + input bool + want string + }{ + {true, "ok"}, + {false, "fail"}, + } + for _, tt := range tests { + got := statusString(tt.input) + if got != tt.want { + t.Errorf("statusString(%v) = %q, want %q", tt.input, got, tt.want) + } + } +} diff --git a/pkg/logger/logger.go b/pkg/logger/logger.go index eeb1436de..1bcc1cec9 100644 --- a/pkg/logger/logger.go +++ b/pkg/logger/logger.go @@ -2,6 +2,7 @@ package logger import ( "fmt" + "io" "os" "path/filepath" "runtime" @@ -100,6 +101,12 @@ func SetConsoleLevel(level LogLevel) { logger = logger.Level(level) } +func DisableConsole() { + mu.Lock() + defer mu.Unlock() + logger = zerolog.New(io.Discard).With().Timestamp().Caller().Logger() +} + func GetLevel() LogLevel { mu.RLock() defer mu.RUnlock() @@ -170,6 +177,22 @@ func DisableFileLogging() { fileLogger = zerolog.Logger{} } +func ConfigureFromEnv() { + if logFile := os.Getenv("PICOCLAW_LOG_FILE"); logFile != "" { + if strings.HasPrefix(logFile, "~/") { + if home := os.Getenv("HOME"); home != "" { + logFile = filepath.Join(home, logFile[2:]) + } + } + + if err := EnableFileLogging(logFile); err != nil { + fmt.Fprintf(os.Stderr, "failed to enable file logging: %v\n", err) + } else { + DisableConsole() + } + } +} + func getCallerSkip() int { for i := 2; i < 15; i++ { pc, file, _, ok := runtime.Caller(i) diff --git a/pkg/logger/logger_test.go b/pkg/logger/logger_test.go index 6ad3a8dd6..1eca72607 100644 --- a/pkg/logger/logger_test.go +++ b/pkg/logger/logger_test.go @@ -4,7 +4,11 @@ import ( "bytes" "encoding/json" "errors" + "fmt" + "os" + "path/filepath" "testing" + "time" "github.com/rs/zerolog" ) @@ -365,3 +369,40 @@ func TestAppendFields_ErrorUsesErrorString(t *testing.T) { t.Fatalf("error field = %#v, want %q", got["error"], "transcription request failed") } } + +func TestDisableConsole(t *testing.T) { + DisableConsole() + Info("this should go to nowhere") +} + +func TestConfigureFromEnv(t *testing.T) { + home := os.Getenv("HOME") + if home == "" { + t.Skip("HOME not set") + } + + tmpFile := "/tmp/picoclaw_test_log_" + fmt.Sprintf("%d", time.Now().UnixNano()) + defer os.Remove(tmpFile) + + os.Setenv("PICOCLAW_LOG_FILE", tmpFile) + defer os.Unsetenv("PICOCLAW_LOG_FILE") + + ConfigureFromEnv() + + if logFile == nil { + t.Error("expected log file to be set") + } + + Info("test message") + + os.Setenv("PICOCLAW_LOG_FILE", "~/test_log") + ConfigureFromEnv() + + expanded := filepath.Join(home, "test_log") + defer os.Remove(expanded) +} + +func TestConfigureFromEnvNoEnv(t *testing.T) { + os.Unsetenv("PICOCLAW_LOG_FILE") + ConfigureFromEnv() +} diff --git a/pkg/migrate/sources/openclaw/common.go b/pkg/migrate/sources/openclaw/common.go index 337c950d0..938f15b80 100644 --- a/pkg/migrate/sources/openclaw/common.go +++ b/pkg/migrate/sources/openclaw/common.go @@ -13,17 +13,16 @@ var migrateableDirs = []string{ } var supportedChannels = map[string]bool{ - "whatsapp": true, - "telegram": true, - "feishu": true, - "discord": true, - "maixcam": true, - "qq": true, - "dingtalk": true, - "slack": true, - "matrix": true, - "line": true, - "onebot": true, - "wecom": true, - "wecom_app": true, + "whatsapp": true, + "telegram": true, + "feishu": true, + "discord": true, + "maixcam": true, + "qq": true, + "dingtalk": true, + "slack": true, + "matrix": true, + "line": true, + "onebot": true, + "wecom": true, } diff --git a/pkg/providers/error_classifier.go b/pkg/providers/error_classifier.go index fd9bf1e81..e7691aa93 100644 --- a/pkg/providers/error_classifier.go +++ b/pkg/providers/error_classifier.go @@ -84,6 +84,15 @@ var ( substr("messages.1.content.1.tool_use.id"), substr("invalid request format"), } + contextOverflowPatterns = []errorPattern{ + rxp(`context[_ ]?length[_ ]?exceeded`), + rxp(`context[_ ]?window[_ ]?exceeded`), + substr("maximum context length"), + substr("token limit"), + substr("too many tokens"), + substr("prompt is too long"), + substr("request too large"), + } imageDimensionPatterns = []errorPattern{ rxp(`image dimensions exceed max`), @@ -201,6 +210,9 @@ func classifyByMessage(msg string) FailoverReason { if matchesAny(msg, formatPatterns) { return FailoverFormat } + if matchesAny(msg, contextOverflowPatterns) { + return FailoverContextOverflow + } return "" } diff --git a/pkg/providers/error_classifier_test.go b/pkg/providers/error_classifier_test.go index 67d9af62b..46b180835 100644 --- a/pkg/providers/error_classifier_test.go +++ b/pkg/providers/error_classifier_test.go @@ -221,6 +221,30 @@ func TestClassifyError_ImageDimensionError(t *testing.T) { } } +func TestClassifyError_ContextOverflowPatterns(t *testing.T) { + patterns := []string{ + "context_length_exceeded", + "context_window_exceeded", + "maximum context length", + "token limit", + "too many tokens", + "prompt is too long", + "request too large", + } + + for _, msg := range patterns { + err := errors.New(msg) + result := ClassifyError(err, "openai", "gpt-4") + if result == nil { + t.Errorf("pattern %q: expected non-nil", msg) + continue + } + if result.Reason != FailoverContextOverflow { + t.Errorf("pattern %q: reason = %q, want context_overflow", msg, result.Reason) + } + } +} + func TestClassifyError_ImageSizeError(t *testing.T) { err := errors.New("image exceeds 20 mb limit") result := ClassifyError(err, "openai", "gpt-4o") @@ -265,6 +289,7 @@ func TestFailoverError_IsRetriable(t *testing.T) { {FailoverTimeout, true}, {FailoverOverloaded, true}, {FailoverFormat, false}, + {FailoverContextOverflow, false}, {FailoverUnknown, true}, } diff --git a/pkg/providers/factory_provider.go b/pkg/providers/factory_provider.go index 1128fc042..962e6ae19 100644 --- a/pkg/providers/factory_provider.go +++ b/pkg/providers/factory_provider.go @@ -158,7 +158,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err "ollama", "moonshot", "shengsuanyun", "deepseek", "cerebras", "vivgrid", "volcengine", "vllm", "qwen", "qwen-intl", "qwen-international", "dashscope-intl", "qwen-us", "dashscope-us", "mistral", "avian", "longcat", "modelscope", "novita", - "coding-plan", "alibaba-coding", "qwen-coding": + "coding-plan", "alibaba-coding", "qwen-coding", "mimo": // 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) @@ -349,6 +349,8 @@ func getDefaultAPIBase(protocol string) string { return "https://api.longcat.chat/openai" case "modelscope": return "https://api-inference.modelscope.cn/v1" + case "mimo": + return "https://api.xiaomimimo.com/v1" default: return "" } diff --git a/pkg/providers/factory_provider_test.go b/pkg/providers/factory_provider_test.go index 2fed18c35..f1fe02cc2 100644 --- a/pkg/providers/factory_provider_test.go +++ b/pkg/providers/factory_provider_test.go @@ -123,6 +123,7 @@ func TestCreateProviderFromConfig_DefaultAPIBase(t *testing.T) { {"ollama", "ollama"}, {"longcat", "longcat"}, {"modelscope", "modelscope"}, + {"mimo", "mimo"}, } for _, tt := range tests { @@ -252,6 +253,35 @@ func TestGetDefaultAPIBase_Novita(t *testing.T) { } } +func TestCreateProviderFromConfig_Mimo(t *testing.T) { + cfg := &config.ModelConfig{ + ModelName: "test-mimo", + Model: "mimo/mimo-v2-pro", + APIBase: "https://api.xiaomimimo.com/v1", + } + cfg.SetAPIKey("test-key") + + 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 != "mimo-v2-pro" { + t.Errorf("modelID = %q, want %q", modelID, "mimo-v2-pro") + } + if _, ok := provider.(*HTTPProvider); !ok { + t.Fatalf("expected *HTTPProvider, got %T", provider) + } +} + +func TestGetDefaultAPIBase_Mimo(t *testing.T) { + if got := getDefaultAPIBase("mimo"); got != "https://api.xiaomimimo.com/v1" { + t.Fatalf("getDefaultAPIBase(%q) = %q, want %q", "mimo", got, "https://api.xiaomimimo.com/v1") + } +} + func TestCreateProviderFromConfig_Anthropic(t *testing.T) { cfg := &config.ModelConfig{ ModelName: "test-anthropic", diff --git a/pkg/providers/types.go b/pkg/providers/types.go index 9a4d126a7..f98ae9243 100644 --- a/pkg/providers/types.go +++ b/pkg/providers/types.go @@ -71,13 +71,14 @@ type NativeSearchCapable interface { type FailoverReason string const ( - FailoverAuth FailoverReason = "auth" - FailoverRateLimit FailoverReason = "rate_limit" - FailoverBilling FailoverReason = "billing" - FailoverTimeout FailoverReason = "timeout" - FailoverFormat FailoverReason = "format" - FailoverOverloaded FailoverReason = "overloaded" - FailoverUnknown FailoverReason = "unknown" + FailoverAuth FailoverReason = "auth" + FailoverRateLimit FailoverReason = "rate_limit" + FailoverBilling FailoverReason = "billing" + FailoverTimeout FailoverReason = "timeout" + FailoverFormat FailoverReason = "format" + FailoverContextOverflow FailoverReason = "context_overflow" + FailoverOverloaded FailoverReason = "overloaded" + FailoverUnknown FailoverReason = "unknown" ) // FailoverError wraps an LLM provider error with classification metadata. @@ -101,7 +102,7 @@ func (e *FailoverError) Unwrap() error { // IsRetriable returns true if this error should trigger fallback to next candidate. // Non-retriable: Format errors (bad request structure, image dimension/size). func (e *FailoverError) IsRetriable() bool { - return e.Reason != FailoverFormat + return e.Reason != FailoverFormat && e.Reason != FailoverContextOverflow } // ModelConfig holds primary model and fallback list. diff --git a/pkg/tools/mcp_tool.go b/pkg/tools/mcp_tool.go index 6e53cf354..5bffb4e89 100644 --- a/pkg/tools/mcp_tool.go +++ b/pkg/tools/mcp_tool.go @@ -5,9 +5,13 @@ import ( "encoding/json" "fmt" "hash/fnv" + "os" "strings" + "time" "github.com/modelcontextprotocol/go-sdk/mcp" + + "github.com/sipeed/picoclaw/pkg/media" ) // MCPManager defines the interface for MCP manager operations @@ -25,6 +29,7 @@ type MCPTool struct { manager MCPManager serverName string tool *mcp.Tool + mediaStore media.MediaStore } // NewMCPTool creates a new MCP tool wrapper @@ -36,6 +41,10 @@ func NewMCPTool(manager MCPManager, serverName string, tool *mcp.Tool) *MCPTool } } +func (t *MCPTool) SetMediaStore(store media.MediaStore) { + t.mediaStore = store +} + // sanitizeIdentifierComponent normalizes a string so it can be safely used // as part of a tool/function identifier for downstream providers. // It: @@ -218,13 +227,7 @@ func (t *MCPTool) Execute(ctx context.Context, args map[string]any) *ToolResult WithError(fmt.Errorf("MCP tool error: %s", errMsg)) } - // Extract text content from result - output := extractContentText(result.Content) - - return &ToolResult{ - ForLLM: output, - IsError: false, - } + return t.normalizeResultContent(ctx, result.Content) } // extractContentText extracts text from MCP content array @@ -233,14 +236,269 @@ func extractContentText(content []mcp.Content) string { for _, c := range content { switch v := c.(type) { case *mcp.TextContent: - parts = append(parts, v.Text) + parts = append(parts, sanitizeToolLLMContent(v.Text)) case *mcp.ImageContent: - // For images, just indicate that an image was returned - parts = append(parts, fmt.Sprintf("[Image: %s]", v.MIMEType)) + parts = append(parts, fmt.Sprintf("[Image: %s]", normalizedMIMEType(v.MIMEType))) + case *mcp.AudioContent: + parts = append(parts, fmt.Sprintf("[Audio: %s]", normalizedMIMEType(v.MIMEType))) + case *mcp.ResourceLink: + parts = append(parts, summarizeResourceLink(v)) + case *mcp.EmbeddedResource: + parts = append(parts, summarizeEmbeddedResource(v)) default: // For other content types, use string representation parts = append(parts, fmt.Sprintf("[Content: %T]", v)) } } - return strings.Join(parts, "\n") + return sanitizeToolLLMContent(strings.Join(parts, "\n")) +} + +func (t *MCPTool) normalizeResultContent(ctx context.Context, content []mcp.Content) *ToolResult { + llmParts := make([]string, 0, len(content)) + mediaRefs := make([]string, 0, len(content)) + + for _, c := range content { + switch v := c.(type) { + case *mcp.TextContent: + text := strings.TrimSpace(sanitizeToolLLMContent(v.Text)) + if text != "" { + llmParts = append(llmParts, text) + } + case *mcp.ImageContent: + ref, note := t.storeBinaryContent( + ctx, + "image", + normalizedMIMEType(v.MIMEType), + v.Data, + v.Annotations, + ) + if ref != "" { + mediaRefs = append(mediaRefs, ref) + } + if note != "" { + llmParts = append(llmParts, note) + } + case *mcp.AudioContent: + ref, note := t.storeBinaryContent( + ctx, + "audio", + normalizedMIMEType(v.MIMEType), + v.Data, + v.Annotations, + ) + if ref != "" { + mediaRefs = append(mediaRefs, ref) + } + if note != "" { + llmParts = append(llmParts, note) + } + case *mcp.ResourceLink: + llmParts = append(llmParts, summarizeResourceLink(v)) + case *mcp.EmbeddedResource: + ref, note := t.storeEmbeddedResource(ctx, v) + if ref != "" { + mediaRefs = append(mediaRefs, ref) + } + if note != "" { + llmParts = append(llmParts, note) + } + default: + llmParts = append(llmParts, fmt.Sprintf("[MCP returned unsupported content type %T]", v)) + } + } + + result := &ToolResult{ + ForLLM: strings.Join(compactStrings(llmParts), "\n"), + Media: mediaRefs, + } + return result +} + +func (t *MCPTool) storeEmbeddedResource(ctx context.Context, content *mcp.EmbeddedResource) (string, string) { + if content == nil || content.Resource == nil { + return "", "[MCP returned an embedded resource without data.]" + } + + resource := content.Resource + if len(resource.Blob) > 0 { + return t.storeBinaryContent( + ctx, + "resource", + normalizedMIMEType(resource.MIMEType), + resource.Blob, + content.Annotations, + ) + } + + if strings.TrimSpace(resource.Text) != "" { + return "", sanitizeToolLLMContent(resource.Text) + } + + return "", summarizeEmbeddedResource(content) +} + +func (t *MCPTool) storeBinaryContent( + ctx context.Context, + kind string, + mimeType string, + data []byte, + annotations *mcp.Annotations, +) (string, string) { + if len(data) == 0 { + return "", fmt.Sprintf("[MCP returned %s content (%s) but it was empty.]", kind, mimeType) + } + if !annotationsAllowUser(annotations) { + return "", fmt.Sprintf( + "[MCP returned %s content (%s) for non-user audience; omitted from model context.]", + kind, + mimeType, + ) + } + if t.mediaStore == nil { + return "", fmt.Sprintf( + "[MCP returned %s content (%s); omitted from model context because media delivery is unavailable.]", + kind, + mimeType, + ) + } + + channel := ToolChannel(ctx) + chatID := ToolChatID(ctx) + if channel == "" || chatID == "" { + return "", fmt.Sprintf( + "[MCP returned %s content (%s); omitted from model context because no target chat was available.]", + kind, + mimeType, + ) + } + + dir := media.TempDir() + if err := os.MkdirAll(dir, 0o700); err != nil { + return "", fmt.Sprintf("[MCP returned %s content (%s) but it could not be stored.]", kind, mimeType) + } + + ext := extensionForMIMEType(mimeType) + tmpFile, err := os.CreateTemp(dir, "mcp-*"+ext) + if err != nil { + return "", fmt.Sprintf("[MCP returned %s content (%s) but it could not be stored.]", kind, mimeType) + } + tmpPath := tmpFile.Name() + if _, err = tmpFile.Write(data); err != nil { + _ = tmpFile.Close() + _ = os.Remove(tmpPath) + return "", fmt.Sprintf("[MCP returned %s content (%s) but it could not be stored.]", kind, mimeType) + } + if err = tmpFile.Close(); err != nil { + _ = os.Remove(tmpPath) + return "", fmt.Sprintf("[MCP returned %s content (%s) but it could not be stored.]", kind, mimeType) + } + + scope := fmt.Sprintf( + "tool:mcp:%s:%s:%s:%d", + sanitizeIdentifierComponent(t.serverName), + channel, + chatID, + time.Now().UnixNano(), + ) + filename := fmt.Sprintf( + "%s_%s%s", + sanitizeIdentifierComponent(t.serverName), + sanitizeIdentifierComponent(t.tool.Name), + ext, + ) + + ref, err := t.mediaStore.Store(tmpPath, media.MediaMeta{ + Filename: filename, + ContentType: mimeType, + Source: fmt.Sprintf( + "tool:mcp:%s:%s", + sanitizeIdentifierComponent(t.serverName), + sanitizeIdentifierComponent(t.tool.Name), + ), + }, scope) + if err != nil { + _ = os.Remove(tmpPath) + return "", fmt.Sprintf( + "[MCP returned %s content (%s) but it could not be registered as media.]", + kind, + mimeType, + ) + } + + return ref, fmt.Sprintf( + "[MCP returned %s content (%s); omitted from model context and stored as a local media artifact.]", + kind, + mimeType, + ) +} + +func summarizeResourceLink(content *mcp.ResourceLink) string { + if content == nil { + return "[MCP returned an empty resource link.]" + } + + parts := []string{"[MCP returned resource link"} + if content.Name != "" { + parts = append(parts, fmt.Sprintf("name=%q", content.Name)) + } + if content.URI != "" { + parts = append(parts, fmt.Sprintf("uri=%q", content.URI)) + } + if content.MIMEType != "" { + parts = append(parts, fmt.Sprintf("mime=%q", content.MIMEType)) + } + if content.Description != "" { + desc := strings.TrimSpace(content.Description) + if len(desc) > 200 { + desc = desc[:200] + "..." + } + parts = append(parts, fmt.Sprintf("description=%q", desc)) + } + return strings.Join(parts, ", ") + "]" +} + +func summarizeEmbeddedResource(content *mcp.EmbeddedResource) string { + if content == nil || content.Resource == nil { + return "[MCP returned an embedded resource.]" + } + + resource := content.Resource + if resource.URI != "" { + return fmt.Sprintf( + "[MCP returned embedded resource %q (%s).]", + resource.URI, + normalizedMIMEType(resource.MIMEType), + ) + } + return fmt.Sprintf("[MCP returned embedded resource (%s).]", normalizedMIMEType(resource.MIMEType)) +} + +func annotationsAllowUser(annotations *mcp.Annotations) bool { + if annotations == nil || len(annotations.Audience) == 0 { + return true + } + for _, audience := range annotations.Audience { + if strings.EqualFold(string(audience), "user") { + return true + } + } + return false +} + +func normalizedMIMEType(mimeType string) string { + if strings.TrimSpace(mimeType) == "" { + return "application/octet-stream" + } + return mimeType +} + +func compactStrings(parts []string) []string { + compact := make([]string, 0, len(parts)) + for _, part := range parts { + if strings.TrimSpace(part) == "" { + continue + } + compact = append(compact, part) + } + return compact } diff --git a/pkg/tools/mcp_tool_test.go b/pkg/tools/mcp_tool_test.go index 95bb0f992..8bbac3bc7 100644 --- a/pkg/tools/mcp_tool_test.go +++ b/pkg/tools/mcp_tool_test.go @@ -3,10 +3,14 @@ package tools import ( "context" "fmt" + "os" + "path/filepath" "strings" "testing" "github.com/modelcontextprotocol/go-sdk/mcp" + + "github.com/sipeed/picoclaw/pkg/media" ) // MockMCPManager is a mock implementation of MCPManager interface for testing @@ -490,3 +494,143 @@ func TestMCPTool_Parameters_MapSchema(t *testing.T) { t.Errorf("Name type should be 'string', got '%v'", nameParam["type"]) } } + +func TestMCPTool_Execute_ImageContentStoredAsMedia(t *testing.T) { + store := media.NewFileMediaStore() + manager := &MockMCPManager{ + callToolFunc: func(ctx context.Context, serverName, toolName string, arguments map[string]any) (*mcp.CallToolResult, error) { + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.ImageContent{ + Data: []byte("fake-image-bytes"), + MIMEType: "image/png", + }, + }, + }, nil + }, + } + + mcpTool := NewMCPTool(manager, "screenshoto", &mcp.Tool{Name: "take_screenshot"}) + mcpTool.SetMediaStore(store) + + result := mcpTool.Execute(WithToolContext(context.Background(), "telegram", "chat-42"), nil) + + if result.IsError { + t.Fatalf("expected success, got %q", result.ForLLM) + } + if len(result.Media) != 1 { + t.Fatalf("expected 1 media ref, got %d", len(result.Media)) + } + if result.ResponseHandled { + t.Fatal("expected MCP image artifact not to mark response as handled") + } + if !strings.Contains(result.ForLLM, "stored as a local media artifact") { + t.Fatalf("expected local media artifact note, got %q", result.ForLLM) + } + + path, meta, err := store.ResolveWithMeta(result.Media[0]) + if err != nil { + t.Fatalf("expected stored media ref to resolve: %v", err) + } + if meta.ContentType != "image/png" { + t.Fatalf("expected image/png content type, got %q", meta.ContentType) + } + if filepath.Ext(path) != ".png" { + t.Fatalf("expected png temp file, got %q", path) + } + data, err := os.ReadFile(path) + if err != nil { + t.Fatalf("expected stored media file to be readable: %v", err) + } + if string(data) != "fake-image-bytes" { + t.Fatalf("expected stored media bytes to match input, got %q", string(data)) + } +} + +func TestMCPTool_Execute_EmbeddedResourceBlobStoredAsMedia(t *testing.T) { + store := media.NewFileMediaStore() + manager := &MockMCPManager{ + callToolFunc: func(ctx context.Context, serverName, toolName string, arguments map[string]any) (*mcp.CallToolResult, error) { + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.EmbeddedResource{ + Resource: &mcp.ResourceContents{ + URI: "file:///tmp/report.png", + MIMEType: "image/png", + Blob: []byte("blob-bytes"), + }, + }, + }, + }, nil + }, + } + + mcpTool := NewMCPTool(manager, "grafana", &mcp.Tool{Name: "get_dashboard_image"}) + mcpTool.SetMediaStore(store) + + result := mcpTool.Execute(WithToolContext(context.Background(), "telegram", "chat-42"), nil) + + if len(result.Media) != 1 { + t.Fatalf("expected embedded resource blob to be stored as media, got %d refs", len(result.Media)) + } + path, _, err := store.ResolveWithMeta(result.Media[0]) + if err != nil { + t.Fatalf("expected stored media ref to resolve: %v", err) + } + data, err := os.ReadFile(path) + if err != nil { + t.Fatalf("expected stored media file to be readable: %v", err) + } + if string(data) != "blob-bytes" { + t.Fatalf("expected stored blob bytes to match input, got %q", string(data)) + } +} + +func TestMCPTool_Execute_RespectsUserAudienceForBinaryContent(t *testing.T) { + store := media.NewFileMediaStore() + manager := &MockMCPManager{ + callToolFunc: func(ctx context.Context, serverName, toolName string, arguments map[string]any) (*mcp.CallToolResult, error) { + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.ImageContent{ + Data: []byte("assistant-only"), + MIMEType: "image/png", + Annotations: &mcp.Annotations{Audience: []mcp.Role{"assistant"}}, + }, + }, + }, nil + }, + } + + mcpTool := NewMCPTool(manager, "screenshoto", &mcp.Tool{Name: "take_screenshot"}) + mcpTool.SetMediaStore(store) + + result := mcpTool.Execute(WithToolContext(context.Background(), "telegram", "chat-42"), nil) + + if len(result.Media) != 0 { + t.Fatalf("expected no media ref for non-user audience, got %d", len(result.Media)) + } + if !strings.Contains(result.ForLLM, "non-user audience") { + t.Fatalf("expected audience note, got %q", result.ForLLM) + } +} + +func TestMCPTool_Execute_LargeBase64TextIsOmittedFromContext(t *testing.T) { + manager := &MockMCPManager{ + callToolFunc: func(ctx context.Context, serverName, toolName string, arguments map[string]any) (*mcp.CallToolResult, error) { + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.TextContent{Text: strings.Repeat("QUJD", 400)}, + }, + }, nil + }, + } + + mcpTool := NewMCPTool(manager, "test_server", &mcp.Tool{Name: "dump_payload"}) + + result := mcpTool.Execute(context.Background(), nil) + + if result.ForLLM != largeBase64OmittedMessage { + t.Fatalf("expected sanitized large base64 note, got %q", result.ForLLM) + } +} diff --git a/pkg/tools/normalization.go b/pkg/tools/normalization.go new file mode 100644 index 000000000..3a76c5d92 --- /dev/null +++ b/pkg/tools/normalization.go @@ -0,0 +1,292 @@ +package tools + +import ( + "encoding/base64" + "fmt" + "mime" + "os" + "path/filepath" + "regexp" + "strings" + "time" + "unicode" + + "github.com/sipeed/picoclaw/pkg/media" +) + +const ( + largeBase64OmittedMessage = "[Tool returned a large base64-like payload; omitted from model context.]" + inlineMediaOmittedMessage = "[Tool returned inline media content; omitted from model context.]" + inlineMediaStoredMessage = "[Tool returned inline media content (%s); omitted from model context and registered as a media attachment.]" +) + +var ( + inlineMarkdownDataURLRe = regexp.MustCompile(`!\[[^\]]*\]\((data:[^)]+)\)`) + inlineRawDataURLRe = regexp.MustCompile(`data:[^;\s]+;base64,[A-Za-z0-9+/=\r\n]+`) +) + +func normalizeToolResult( + result *ToolResult, + toolName string, + store media.MediaStore, + channel string, + chatID string, +) *ToolResult { + if result == nil { + return nil + } + + notes := make([]string, 0, 2) + seen := make(map[string]struct{}) + + if store != nil && channel != "" && chatID != "" { + var refs []string + var extractedNotes []string + + result.ForLLM, refs, extractedNotes = extractInlineMediaRefs( + result.ForLLM, + toolName, + store, + channel, + chatID, + seen, + ) + result.Media = append(result.Media, refs...) + notes = append(notes, extractedNotes...) + + result.ForUser, refs, extractedNotes = extractInlineMediaRefs( + result.ForUser, + toolName, + store, + channel, + chatID, + seen, + ) + result.Media = append(result.Media, refs...) + notes = append(notes, extractedNotes...) + } + + result.ForLLM = sanitizeToolLLMContent(result.ForLLM) + + if len(result.Media) > 0 && len(notes) > 0 { + if strings.TrimSpace(result.ForLLM) == "" { + result.ForLLM = strings.Join(notes, "\n") + } else { + result.ForLLM = strings.TrimSpace(result.ForLLM) + "\n" + strings.Join(notes, "\n") + } + } + if len(result.Media) > 0 && strings.TrimSpace(result.ForLLM) == "" { + result.ForLLM = "[Tool returned media content; omitted from model context and registered as a media attachment.]" + } + + return result +} + +func sanitizeToolLLMContent(text string) string { + trimmed := strings.TrimSpace(text) + if trimmed == "" { + return text + } + if inlineMarkdownDataURLRe.MatchString(trimmed) || inlineRawDataURLRe.MatchString(trimmed) { + cleaned := inlineMarkdownDataURLRe.ReplaceAllString(trimmed, "") + cleaned = inlineRawDataURLRe.ReplaceAllString(cleaned, "") + cleaned = strings.TrimSpace(cleaned) + if cleaned == "" { + return inlineMediaOmittedMessage + } + return cleaned + "\n" + inlineMediaOmittedMessage + } + if looksLikeLargeBase64Payload(trimmed) { + return largeBase64OmittedMessage + } + return text +} + +func looksLikeLargeBase64Payload(text string) bool { + trimmed := strings.TrimSpace(text) + if len(trimmed) < 1024 { + return false + } + + nonSpace := 0 + base64Like := 0 + spaceCount := 0 + + for _, r := range trimmed { + if unicode.IsSpace(r) { + spaceCount++ + continue + } + nonSpace++ + if (r >= 'A' && r <= 'Z') || + (r >= 'a' && r <= 'z') || + (r >= '0' && r <= '9') || + r == '+' || r == '/' || r == '=' { + base64Like++ + } + } + + if nonSpace == 0 { + return false + } + + ratio := float64(base64Like) / float64(nonSpace) + return ratio >= 0.97 && spaceCount <= len(trimmed)/128 +} + +func extractInlineMediaRefs( + text string, + toolName string, + store media.MediaStore, + channel string, + chatID string, + seen map[string]struct{}, +) (cleaned string, refs []string, notes []string) { + cleaned = text + + matches := inlineMarkdownDataURLRe.FindAllStringSubmatch(cleaned, -1) + for _, match := range matches { + if len(match) < 2 { + continue + } + dataURL := match[1] + ref, note := storeInlineDataURL(toolName, store, channel, chatID, dataURL, seen) + if ref != "" { + refs = append(refs, ref) + } + if note != "" { + notes = append(notes, note) + } + cleaned = strings.ReplaceAll(cleaned, match[0], "") + } + + rawMatches := inlineRawDataURLRe.FindAllString(cleaned, -1) + for _, dataURL := range rawMatches { + ref, note := storeInlineDataURL(toolName, store, channel, chatID, dataURL, seen) + if ref != "" { + refs = append(refs, ref) + } + if note != "" { + notes = append(notes, note) + } + cleaned = strings.ReplaceAll(cleaned, dataURL, "") + } + + return strings.TrimSpace(cleaned), refs, notes +} + +func storeInlineDataURL( + toolName string, + store media.MediaStore, + channel string, + chatID string, + dataURL string, + seen map[string]struct{}, +) (ref string, note string) { + dataURL = strings.TrimSpace(dataURL) + if _, ok := seen[dataURL]; ok { + return "", "" + } + seen[dataURL] = struct{}{} + + if !strings.HasPrefix(strings.ToLower(dataURL), "data:") { + return "", "" + } + + comma := strings.IndexByte(dataURL, ',') + if comma <= 5 { + return "", "[Tool returned inline media content that could not be parsed.]" + } + + metaPart := dataURL[:comma] + payload := dataURL[comma+1:] + if !strings.Contains(strings.ToLower(metaPart), ";base64") { + return "", "[Tool returned inline media content that was not base64-encoded.]" + } + + mimeType := strings.TrimSpace(strings.TrimPrefix(metaPart, "data:")) + if semi := strings.IndexByte(mimeType, ';'); semi >= 0 { + mimeType = mimeType[:semi] + } + if mimeType == "" { + mimeType = "application/octet-stream" + } + + payload = strings.NewReplacer("\n", "", "\r", "", "\t", "", " ", "").Replace(payload) + decoded, err := base64.StdEncoding.DecodeString(payload) + if err != nil { + return "", fmt.Sprintf("[Tool returned inline media content (%s) that could not be decoded.]", mimeType) + } + + dir := media.TempDir() + if err = os.MkdirAll(dir, 0o700); err != nil { + return "", fmt.Sprintf("[Tool returned inline media content (%s) but it could not be stored.]", mimeType) + } + + ext := extensionForMIMEType(mimeType) + tmpFile, err := os.CreateTemp(dir, "tool-inline-*"+ext) + if err != nil { + return "", fmt.Sprintf("[Tool returned inline media content (%s) but it could not be stored.]", mimeType) + } + tmpPath := tmpFile.Name() + if _, err = tmpFile.Write(decoded); err != nil { + tmpFile.Close() + _ = os.Remove(tmpPath) + return "", fmt.Sprintf("[Tool returned inline media content (%s) but it could not be stored.]", mimeType) + } + if err = tmpFile.Close(); err != nil { + _ = os.Remove(tmpPath) + return "", fmt.Sprintf("[Tool returned inline media content (%s) but it could not be stored.]", mimeType) + } + + filename := sanitizeIdentifierComponent(toolName) + ext + scope := fmt.Sprintf( + "tool:inline:%s:%s:%s:%d", + sanitizeIdentifierComponent(toolName), + channel, + chatID, + time.Now().UnixNano(), + ) + + ref, err = store.Store(tmpPath, media.MediaMeta{ + Filename: filename, + ContentType: mimeType, + Source: fmt.Sprintf("tool:inline:%s", sanitizeIdentifierComponent(toolName)), + }, scope) + if err != nil { + _ = os.Remove(tmpPath) + return "", fmt.Sprintf("[Tool returned inline media content (%s) but it could not be registered.]", mimeType) + } + + return ref, fmt.Sprintf(inlineMediaStoredMessage, mimeType) +} + +func extensionForMIMEType(mimeType string) string { + if mimeType == "" { + return ".bin" + } + if exts, err := mime.ExtensionsByType(mimeType); err == nil && len(exts) > 0 { + return exts[0] + } + + switch strings.ToLower(mimeType) { + case "image/jpeg": + return ".jpg" + case "image/png": + return ".png" + case "image/gif": + return ".gif" + case "image/webp": + return ".webp" + case "audio/wav", "audio/x-wav": + return ".wav" + case "audio/mpeg": + return ".mp3" + case "audio/ogg": + return ".ogg" + case "video/mp4": + return ".mp4" + default: + return filepath.Ext(mimeType) + } +} diff --git a/pkg/tools/registry.go b/pkg/tools/registry.go index 2c634e673..56af8d695 100644 --- a/pkg/tools/registry.go +++ b/pkg/tools/registry.go @@ -9,6 +9,7 @@ import ( "time" "github.com/sipeed/picoclaw/pkg/logger" + "github.com/sipeed/picoclaw/pkg/media" "github.com/sipeed/picoclaw/pkg/providers" ) @@ -19,9 +20,14 @@ type ToolEntry struct { } type ToolRegistry struct { - tools map[string]*ToolEntry - mu sync.RWMutex - version atomic.Uint64 // incremented on Register/RegisterHidden for cache invalidation + tools map[string]*ToolEntry + mu sync.RWMutex + version atomic.Uint64 // incremented on Register/RegisterHidden for cache invalidation + mediaStore media.MediaStore +} + +type mediaStoreAware interface { + SetMediaStore(store media.MediaStore) } func NewToolRegistry() *ToolRegistry { @@ -43,6 +49,9 @@ func (r *ToolRegistry) Register(tool Tool) { IsCore: true, TTL: 0, // Core tools do not use TTL } + if aware, ok := tool.(mediaStoreAware); ok && r.mediaStore != nil { + aware.SetMediaStore(r.mediaStore) + } r.version.Add(1) logger.DebugCF("tools", "Registered core tool", map[string]any{"name": name}) } @@ -61,10 +70,27 @@ func (r *ToolRegistry) RegisterHidden(tool Tool) { IsCore: false, TTL: 0, } + if aware, ok := tool.(mediaStoreAware); ok && r.mediaStore != nil { + aware.SetMediaStore(r.mediaStore) + } r.version.Add(1) logger.DebugCF("tools", "Registered hidden tool", map[string]any{"name": name}) } +// SetMediaStore injects a MediaStore into all registered tools that can +// consume it, and remembers it for future registrations. +func (r *ToolRegistry) SetMediaStore(store media.MediaStore) { + r.mu.Lock() + defer r.mu.Unlock() + + r.mediaStore = store + for _, entry := range r.tools { + if aware, ok := entry.Tool.(mediaStoreAware); ok { + aware.SetMediaStore(store) + } + } +} + // PromoteTools atomically sets the TTL for multiple non-core tools. // This prevents a concurrent TickTTL from decrementing between promotions. func (r *ToolRegistry) PromoteTools(names []string, ttl int) { @@ -238,6 +264,8 @@ func (r *ToolRegistry) ExecuteWithContext( } } + result = normalizeToolResult(result, name, r.mediaStore, channel, chatID) + duration := time.Since(start) // Log based on result type @@ -259,7 +287,7 @@ func (r *ToolRegistry) ExecuteWithContext( map[string]any{ "tool": name, "duration_ms": duration.Milliseconds(), - "result_length": len(result.ForLLM), + "result_length": len(result.ContentForLLM()), }) } @@ -354,7 +382,8 @@ func (r *ToolRegistry) Clone() *ToolRegistry { r.mu.RLock() defer r.mu.RUnlock() clone := &ToolRegistry{ - tools: make(map[string]*ToolEntry, len(r.tools)), + tools: make(map[string]*ToolEntry, len(r.tools)), + mediaStore: r.mediaStore, } for name, entry := range r.tools { clone.tools[name] = &ToolEntry{ diff --git a/pkg/tools/registry_test.go b/pkg/tools/registry_test.go index 967758dfa..db52749f6 100644 --- a/pkg/tools/registry_test.go +++ b/pkg/tools/registry_test.go @@ -3,10 +3,13 @@ package tools import ( "context" "errors" + "os" + "path/filepath" "strings" "sync" "testing" + "github.com/sipeed/picoclaw/pkg/media" "github.com/sipeed/picoclaw/pkg/providers" ) @@ -46,6 +49,15 @@ func (m *mockAsyncRegistryTool) ExecuteAsync(_ context.Context, args map[string] return m.result } +type mockMediaStoreAwareTool struct { + mockRegistryTool + store media.MediaStore +} + +func (m *mockMediaStoreAwareTool) SetMediaStore(store media.MediaStore) { + m.store = store +} + // --- helpers --- func newMockTool(name, desc string) *mockRegistryTool { @@ -621,3 +633,102 @@ func TestToolRegistry_Execute_PanicDoesNotAffectOtherTools(t *testing.T) { t.Errorf("expected 'success', got %q", result2.ForLLM) } } + +func TestToolRegistry_SetMediaStore_PropagatesToExistingAndNewTools(t *testing.T) { + r := NewToolRegistry() + store := media.NewFileMediaStore() + + existing := &mockMediaStoreAwareTool{ + mockRegistryTool: *newMockTool("existing", "existing tool"), + } + r.Register(existing) + + r.SetMediaStore(store) + if existing.store != store { + t.Fatal("expected existing tool to receive media store") + } + + later := &mockMediaStoreAwareTool{ + mockRegistryTool: *newMockTool("later", "later tool"), + } + r.Register(later) + + if later.store != store { + t.Fatal("expected newly registered tool to inherit media store") + } +} + +func TestToolRegistry_ExecuteWithContext_SanitizesLargeBase64Payload(t *testing.T) { + r := NewToolRegistry() + payload := strings.Repeat("QUJD", 400) + r.Register(&mockRegistryTool{ + name: "base64_tool", + desc: "returns huge base64", + params: map[string]any{}, + result: SilentResult(payload), + }) + + result := r.ExecuteWithContext(context.Background(), "base64_tool", nil, "telegram", "chat-1", nil) + + if result.ForLLM != largeBase64OmittedMessage { + t.Fatalf("expected sanitized payload, got %q", result.ForLLM) + } +} + +func TestToolRegistry_ExecuteWithContext_ExtractsInlineMediaDataURL(t *testing.T) { + r := NewToolRegistry() + store := media.NewFileMediaStore() + r.SetMediaStore(store) + + payload := "![screenshot](data:image/png;base64,aGVsbG8=)" + r.Register(&mockRegistryTool{ + name: "inline_media_tool", + desc: "returns inline data url", + params: map[string]any{}, + result: SilentResult(payload), + }) + + result := r.ExecuteWithContext(context.Background(), "inline_media_tool", nil, "telegram", "chat-42", nil) + + if len(result.Media) != 1 { + t.Fatalf("expected 1 media ref, got %d", len(result.Media)) + } + if strings.Contains(result.ForLLM, "data:image/png;base64") { + t.Fatalf("expected inline data URL to be stripped from ForLLM, got %q", result.ForLLM) + } + if !strings.Contains(result.ForLLM, "registered as a media attachment") { + t.Fatalf("expected delivery note in ForLLM, got %q", result.ForLLM) + } + + path, err := store.Resolve(result.Media[0]) + if err != nil { + t.Fatalf("expected stored media ref to resolve: %v", err) + } + if _, err := os.Stat(path); err != nil { + t.Fatalf("expected stored media file to exist: %v", err) + } + if filepath.Ext(path) != ".png" { + t.Fatalf("expected stored inline media to use png extension, got %q", path) + } +} + +func TestToolRegistry_ExecuteWithContext_SanitizesInlineMediaWithoutStore(t *testing.T) { + r := NewToolRegistry() + + payload := "before ![img](data:image/png;base64,aGVsbG8=) after" + r.Register(&mockRegistryTool{ + name: "inline_media_no_store", + desc: "returns inline data url without store", + params: map[string]any{}, + result: SilentResult(payload), + }) + + result := r.ExecuteWithContext(context.Background(), "inline_media_no_store", nil, "telegram", "chat-42", nil) + + if strings.Contains(result.ForLLM, "data:image/png;base64") { + t.Fatalf("expected inline data URL to be removed from ForLLM, got %q", result.ForLLM) + } + if !strings.Contains(result.ForLLM, inlineMediaOmittedMessage) { + t.Fatalf("expected inline media omission note, got %q", result.ForLLM) + } +} diff --git a/pkg/tools/result.go b/pkg/tools/result.go index bf34b7bc6..c81213125 100644 --- a/pkg/tools/result.go +++ b/pkg/tools/result.go @@ -2,10 +2,16 @@ package tools import ( "encoding/json" + "strings" "github.com/sipeed/picoclaw/pkg/providers" ) +const ( + handledToolLLMNote = "The requested output has already been delivered to the user in the current chat. Do not call send_file or any other delivery tool again. If you reply, provide only a brief confirmation." + artifactPathsLLMNote = "Use `send_file` with one of these paths to send it to the user, or use file/exec tools to save it inside the workspace if requested." +) + // ToolResult represents the structured return value from tool execution. // It provides clear semantics for different types of results and supports // async operations, user-facing messages, and error handling. @@ -43,6 +49,48 @@ type ToolResult struct { // Only populated by SubTurn executions; used by evaluator_optimizer // to carry stateful worker context across evaluation iterations. Messages []providers.Message `json:"-"` + + // ArtifactTags exposes local artifact paths back to the LLM in a structured + // form, e.g. "[file:/tmp/example.png]". This is used when a tool produced a + // reusable local artifact but did not deliver it to the user yet. + ArtifactTags []string `json:"artifact_tags,omitempty"` + + // ResponseHandled indicates that this tool execution already satisfied the + // user's request at the channel/output level, so the agent loop can stop + // without a follow-up assistant response. + ResponseHandled bool `json:"response_handled,omitempty"` +} + +// ContentForLLM returns the normalized textual content to append to the +// conversation after a tool call. Errors fall back to Err when ForLLM is empty. +func (tr *ToolResult) ContentForLLM() string { + if tr == nil { + return "" + } + content := tr.ForLLM + if content == "" && tr.Err != nil { + content = tr.Err.Error() + } + if tr.ResponseHandled { + if content == "" { + return handledToolLLMNote + } + if !strings.Contains(content, handledToolLLMNote) { + content += "\n" + handledToolLLMNote + } + } + if len(tr.ArtifactTags) > 0 { + artifactNote := "Local artifact paths: " + strings.Join(tr.ArtifactTags, " ") + "\n" + artifactPathsLLMNote + if content == "" { + content = artifactNote + } else if !strings.Contains(content, artifactNote) { + content += "\n" + artifactNote + } + } + if content != "" { + return content + } + return "" } // NewToolResult creates a basic ToolResult with content for the LLM. @@ -167,3 +215,9 @@ func (tr *ToolResult) WithError(err error) *ToolResult { tr.Err = err return tr } + +// WithResponseHandled marks the tool result as already delivered to the user. +func (tr *ToolResult) WithResponseHandled() *ToolResult { + tr.ResponseHandled = true + return tr +} diff --git a/pkg/tools/result_test.go b/pkg/tools/result_test.go index a234e33f3..5f08cb4fa 100644 --- a/pkg/tools/result_test.go +++ b/pkg/tools/result_test.go @@ -3,6 +3,7 @@ package tools import ( "encoding/json" "errors" + "strings" "testing" ) @@ -227,3 +228,41 @@ func TestToolResultJSONStructure(t *testing.T) { t.Errorf("Expected silent false, got %v", parsed["silent"]) } } + +func TestToolResultContentForLLM_AppendsHandledDeliveryNote(t *testing.T) { + result := MediaResult("Screenshot attached.", []string{"media://example"}).WithResponseHandled() + + content := result.ContentForLLM() + if !strings.Contains(content, "Screenshot attached.") { + t.Fatalf("expected original content in ContentForLLM, got %q", content) + } + if !strings.Contains(content, handledToolLLMNote) { + t.Fatalf("expected handled delivery note in ContentForLLM, got %q", content) + } +} + +func TestToolResultContentForLLM_UsesHandledDeliveryNoteWhenEmpty(t *testing.T) { + result := (&ToolResult{}).WithResponseHandled() + + if got := result.ContentForLLM(); got != handledToolLLMNote { + t.Fatalf("ContentForLLM() = %q, want %q", got, handledToolLLMNote) + } +} + +func TestToolResultContentForLLM_AppendsArtifactPaths(t *testing.T) { + result := &ToolResult{ + ForLLM: "Artifact created.", + ArtifactTags: []string{"[file:/tmp/example.png]"}, + } + + content := result.ContentForLLM() + if !strings.Contains(content, "Artifact created.") { + t.Fatalf("expected original content in ContentForLLM, got %q", content) + } + if !strings.Contains(content, "Local artifact paths: [file:/tmp/example.png]") { + t.Fatalf("expected artifact path note in ContentForLLM, got %q", content) + } + if !strings.Contains(content, artifactPathsLLMNote) { + t.Fatalf("expected artifact guidance note in ContentForLLM, got %q", content) + } +} diff --git a/pkg/tools/send_file.go b/pkg/tools/send_file.go index 57b99a845..44198381e 100644 --- a/pkg/tools/send_file.go +++ b/pkg/tools/send_file.go @@ -142,7 +142,7 @@ func (t *SendFileTool) Execute(ctx context.Context, args map[string]any) *ToolRe return ErrorResult(fmt.Sprintf("failed to register media: %v", err)) } - return MediaResult(fmt.Sprintf("File %q sent to user", filename), []string{ref}) + return MediaResult(fmt.Sprintf("File %q sent to user", filename), []string{ref}).WithResponseHandled() } // detectMediaType determines the MIME type of a file. diff --git a/pkg/tools/send_file_test.go b/pkg/tools/send_file_test.go index 0a99e8028..f36baf7d0 100644 --- a/pkg/tools/send_file_test.go +++ b/pkg/tools/send_file_test.go @@ -104,6 +104,9 @@ func TestSendFileTool_Success(t *testing.T) { if result.Media[0][:8] != "media://" { t.Errorf("expected media:// ref, got %q", result.Media[0]) } + if !result.ResponseHandled { + t.Fatal("expected send_file success to mark response handled") + } _, meta, err := store.ResolveWithMeta(result.Media[0]) if err != nil { diff --git a/pkg/tools/session.go b/pkg/tools/session.go new file mode 100644 index 000000000..141dd4b5e --- /dev/null +++ b/pkg/tools/session.go @@ -0,0 +1,252 @@ +package tools + +import ( + "bytes" + "errors" + "io" + "os" + "sync" + "time" + + "github.com/google/uuid" +) + +const maxOutputBufferSize = 1 * 1024 * 1024 // 1MB + +const outputTruncateMarker = "\n... [output truncated, exceeded 1MB]\n" + +// PtyKeyMode represents arrow key encoding mode for PTY sessions. +// Programs send smkx/rmkx sequences to switch between CSI and SS3 modes. +type PtyKeyMode uint8 + +const ( + PtyKeyModeCSI PtyKeyMode = iota // triggered by rmkx (\x1b[?1l) + PtyKeyModeSS3 // triggered by smkx (\x1b[?1h) +) + +const PtyKeyModeNotFound PtyKeyMode = 255 + +var ( + ErrSessionNotFound = errors.New("session not found") + ErrSessionDone = errors.New("session already completed") + ErrPTYNotSupported = errors.New("PTY is not supported on this platform") + ErrNoStdin = errors.New("no stdin available") +) + +type ProcessSession struct { + mu sync.Mutex + ID string + PID int + Command string + PTY bool + Background bool + StartTime int64 + ExitCode int + Status string + stdinWriter io.Writer + stdoutPipe io.Reader + outputBuffer *bytes.Buffer + outputTruncated bool + ptyMaster *os.File + + // ptyKeyMode tracks arrow key encoding mode (CSI vs SS3) + ptyKeyMode PtyKeyMode +} + +func (s *ProcessSession) IsDone() bool { + s.mu.Lock() + defer s.mu.Unlock() + return s.Status == "done" || s.Status == "exited" +} + +func (s *ProcessSession) GetPtyKeyMode() PtyKeyMode { + s.mu.Lock() + defer s.mu.Unlock() + return s.ptyKeyMode +} + +func (s *ProcessSession) SetPtyKeyMode(mode PtyKeyMode) { + s.mu.Lock() + defer s.mu.Unlock() + s.ptyKeyMode = mode +} + +func (s *ProcessSession) GetStatus() string { + s.mu.Lock() + defer s.mu.Unlock() + return s.Status +} + +func (s *ProcessSession) SetStatus(status string) { + s.mu.Lock() + defer s.mu.Unlock() + s.Status = status +} + +func (s *ProcessSession) GetExitCode() int { + s.mu.Lock() + defer s.mu.Unlock() + return s.ExitCode +} + +func (s *ProcessSession) SetExitCode(code int) { + s.mu.Lock() + defer s.mu.Unlock() + s.ExitCode = code +} + +func (s *ProcessSession) killProcess() error { + s.mu.Lock() + defer s.mu.Unlock() + + if s.Status != "running" { + return ErrSessionDone + } + + pid := s.PID + if pid <= 0 { + return ErrSessionNotFound + } + + if err := killProcessGroup(pid); err != nil { + return err + } + + s.Status = "done" + s.ExitCode = -1 + return nil +} + +func (s *ProcessSession) Kill() error { + return s.killProcess() +} + +func (s *ProcessSession) Write(data string) error { + s.mu.Lock() + defer s.mu.Unlock() + + if s.Status != "running" { + return ErrSessionDone + } + + var writer io.Writer + if s.PTY && s.ptyMaster != nil { + writer = s.ptyMaster + } else if s.stdinWriter != nil { + writer = s.stdinWriter + } else { + return ErrNoStdin + } + + _, err := writer.Write([]byte(data)) + return err +} + +func (s *ProcessSession) Read() string { + s.mu.Lock() + defer s.mu.Unlock() + + if s.outputBuffer.Len() == 0 { + return "" + } + + data := s.outputBuffer.String() + s.outputBuffer.Reset() + return data +} + +func (s *ProcessSession) ToSessionInfo() SessionInfo { + s.mu.Lock() + defer s.mu.Unlock() + + return SessionInfo{ + ID: s.ID, + Command: s.Command, + Status: s.Status, + PID: s.PID, + StartedAt: s.StartTime, + } +} + +type SessionManager struct { + mu sync.RWMutex + sessions map[string]*ProcessSession +} + +func NewSessionManager() *SessionManager { + sm := &SessionManager{ + sessions: make(map[string]*ProcessSession), + } + + // Start cleaner goroutine - runs every 5 minutes, cleans up sessions done for >30 minutes + go func() { + ticker := time.NewTicker(5 * time.Minute) + defer ticker.Stop() + for range ticker.C { + sm.cleanupOldSessions() + } + }() + + return sm +} + +// cleanupOldSessions removes sessions that are done and older than 30 minutes +func (sm *SessionManager) cleanupOldSessions() { + sm.mu.Lock() + defer sm.mu.Unlock() + + cutoff := time.Now().Add(-30 * time.Minute) + for id, session := range sm.sessions { + if session.IsDone() && session.StartTime < cutoff.Unix() { + delete(sm.sessions, id) + } + } +} + +func (sm *SessionManager) Add(session *ProcessSession) { + sm.mu.Lock() + defer sm.mu.Unlock() + sm.sessions[session.ID] = session +} + +func (sm *SessionManager) Get(sessionID string) (*ProcessSession, error) { + sm.mu.RLock() + defer sm.mu.RUnlock() + + session, ok := sm.sessions[sessionID] + if !ok { + return nil, ErrSessionNotFound + } + + return session, nil +} + +func (sm *SessionManager) Remove(sessionID string) { + sm.mu.Lock() + defer sm.mu.Unlock() + delete(sm.sessions, sessionID) +} + +func (sm *SessionManager) List() []SessionInfo { + sm.mu.RLock() + defer sm.mu.RUnlock() + + result := make([]SessionInfo, 0, len(sm.sessions)) + for _, session := range sm.sessions { + result = append(result, session.ToSessionInfo()) + } + + return result +} + +func generateSessionID() string { + return uuid.New().String()[:8] +} + +type SessionInfo struct { + ID string `json:"id"` + Command string `json:"command"` + Status string `json:"status"` + PID int `json:"pid"` + StartedAt int64 `json:"startedAt"` +} diff --git a/pkg/tools/session_process_unix.go b/pkg/tools/session_process_unix.go new file mode 100644 index 000000000..2fe30166e --- /dev/null +++ b/pkg/tools/session_process_unix.go @@ -0,0 +1,14 @@ +//go:build !windows + +package tools + +import ( + "syscall" +) + +func killProcessGroup(pid int) error { + if err := syscall.Kill(-pid, syscall.SIGKILL); err != nil { + _ = syscall.Kill(pid, syscall.SIGKILL) + } + return nil +} diff --git a/pkg/tools/session_process_windows.go b/pkg/tools/session_process_windows.go new file mode 100644 index 000000000..7cf558954 --- /dev/null +++ b/pkg/tools/session_process_windows.go @@ -0,0 +1,13 @@ +//go:build windows + +package tools + +import ( + "os/exec" + "strconv" +) + +func killProcessGroup(pid int) error { + _ = exec.Command("taskkill", "/T", "/F", "/PID", strconv.Itoa(pid)).Run() + return nil +} diff --git a/pkg/tools/session_test.go b/pkg/tools/session_test.go new file mode 100644 index 000000000..6cfe72a10 --- /dev/null +++ b/pkg/tools/session_test.go @@ -0,0 +1,99 @@ +package tools + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestSessionManager_AddGet(t *testing.T) { + sm := NewSessionManager() + session := &ProcessSession{ + ID: "test-1", + Command: "echo hello", + Status: "running", + StartTime: 1000, + } + + sm.Add(session) + + got, err := sm.Get("test-1") + require.NoError(t, err) + require.Equal(t, "test-1", got.ID) +} + +func TestSessionManager_Remove(t *testing.T) { + sm := NewSessionManager() + session := &ProcessSession{ + ID: "test-1", + Command: "echo hello", + Status: "running", + StartTime: 1000, + } + sm.Add(session) + sm.Remove("test-1") + + _, err := sm.Get("test-1") + require.ErrorIs(t, err, ErrSessionNotFound) +} + +func TestSessionManager_List(t *testing.T) { + sm := NewSessionManager() + sm.Add(&ProcessSession{ + ID: "test-1", + Command: "echo hello", + Status: "running", + StartTime: 1000, + }) + sm.Add(&ProcessSession{ + ID: "test-2", + Command: "echo world", + Status: "running", + StartTime: 1001, + }) + sm.Add(&ProcessSession{ + ID: "test-3", + Command: "echo done", + Status: "done", + StartTime: 1002, + }) + + sessions := sm.List() + require.Len(t, sessions, 3) + + ids := make(map[string]bool) + for _, s := range sessions { + ids[s.ID] = true + } + require.True(t, ids["test-1"]) + require.True(t, ids["test-2"]) + require.True(t, ids["test-3"]) +} + +func TestProcessSession_IsDone(t *testing.T) { + session := &ProcessSession{Status: "running"} + require.False(t, session.IsDone()) + + session.Status = "done" + require.True(t, session.IsDone()) + + session.Status = "exited" + require.True(t, session.IsDone()) +} + +func TestProcessSession_ToSessionInfo(t *testing.T) { + session := &ProcessSession{ + ID: "test-1", + PID: 12345, + Command: "echo hello", + Status: "running", + StartTime: 1000, + } + + info := session.ToSessionInfo() + require.Equal(t, "test-1", info.ID) + require.Equal(t, "echo hello", info.Command) + require.Equal(t, "running", info.Status) + require.Equal(t, 12345, info.PID) + require.Equal(t, int64(1000), info.StartedAt) +} diff --git a/pkg/tools/shell.go b/pkg/tools/shell.go index 78ad2b26d..6ee1cb993 100644 --- a/pkg/tools/shell.go +++ b/pkg/tools/shell.go @@ -3,20 +3,36 @@ package tools import ( "bytes" "context" + "encoding/json" "errors" "fmt" + "io" "os" "os/exec" "path/filepath" "regexp" "runtime" "strings" + "sync" "time" + "github.com/creack/pty" + "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/constants" ) +var ( + globalSessionManager = NewSessionManager() + sessionManagerMu sync.RWMutex +) + +func getSessionManager() *SessionManager { + sessionManagerMu.RLock() + defer sessionManagerMu.RUnlock() + return globalSessionManager +} + type ExecTool struct { workingDir string timeout time.Duration @@ -26,6 +42,7 @@ type ExecTool struct { allowedPathPatterns []*regexp.Regexp restrictToWorkspace bool allowRemote bool + sessionManager *SessionManager } var ( @@ -145,7 +162,7 @@ func NewExecToolWithConfig( denyPatterns = append(denyPatterns, defaultDenyPatterns...) } - timeout := 60 * time.Second + var timeout time.Duration if config != nil && config.Tools.Exec.TimeoutSeconds > 0 { timeout = time.Duration(config.Tools.Exec.TimeoutSeconds) * time.Second } @@ -159,6 +176,7 @@ func NewExecToolWithConfig( allowedPathPatterns: allowedPathPatterns, restrictToWorkspace: restrict, allowRemote: allowRemote, + sessionManager: getSessionManager(), }, nil } @@ -167,27 +185,82 @@ func (t *ExecTool) Name() string { } func (t *ExecTool) Description() string { - return "Execute a shell command and return its output. Use with caution." + return `Execute shell commands. Use background=true for long-running commands (returns sessionId). Use pty=true for interactive commands (can combine with background=true). Use poll/read/write/send-keys/kill with sessionId to manage background sessions. Sessions auto-cleanup 30 minutes after process exits; use kill to terminate early. Output buffer limit: 1MB.` } func (t *ExecTool) Parameters() map[string]any { return map[string]any{ "type": "object", "properties": map[string]any{ + "action": map[string]any{ + "type": "string", + "enum": []string{"run", "list", "poll", "read", "write", "kill", "send-keys"}, + "description": "Action: run (execute command), list (show sessions), poll (check status), read (get output), write (send input), kill (terminate), send-keys (send keys to PTY)", + }, "command": map[string]any{ "type": "string", - "description": "The shell command to execute", + "description": "Shell command to execute (required for run)", }, - "working_dir": map[string]any{ + "sessionId": map[string]any{ "type": "string", - "description": "Optional working directory for the command", + "description": "Session ID (required for poll/read/write/kill/send-keys)", + }, + "keys": map[string]any{ + "type": "string", + "description": "Key names for send-keys: up, down, left, right, enter, tab, escape, backspace, ctrl-c, ctrl-d, home, end, pageup, pagedown, f1-f12", + }, + "data": map[string]any{ + "type": "string", + "description": "Data to write to stdin (required for write)", + }, + "background": map[string]any{ + "type": "string", + "description": "Run in background immediately", + }, + "pty": map[string]any{ + "type": "string", + "description": "Run in a pseudo-terminal (PTY) when available", + }, + "cwd": map[string]any{ + "type": "string", + "description": "Working directory for the command", + }, + "timeout": map[string]any{ + "type": "integer", + "description": "Timeout in seconds (0 = no timeout)", }, }, - "required": []string{"command"}, + "required": []string{"action"}, } } func (t *ExecTool) Execute(ctx context.Context, args map[string]any) *ToolResult { + action, _ := args["action"].(string) + if action == "" { + return ErrorResult("action is required") + } + + switch action { + case "run": + return t.executeRun(ctx, args) + case "list": + return t.executeList() + case "poll": + return t.executePoll(args) + case "read": + return t.executeRead(args) + case "write": + return t.executeWrite(args) + case "kill": + return t.executeKill(args) + case "send-keys": + return t.executeSendKeys(args) + default: + return ErrorResult(fmt.Sprintf("unknown action: %s", action)) + } +} + +func (t *ExecTool) executeRun(ctx context.Context, args map[string]any) *ToolResult { command, ok := args["command"].(string) if !ok { return ErrorResult("command is required") @@ -206,8 +279,26 @@ func (t *ExecTool) Execute(ctx context.Context, args map[string]any) *ToolResult } } + getBoolArg := func(key string) bool { + switch v := args[key].(type) { + case bool: + return v + case string: + return v == "true" + } + return false + } + isPty := getBoolArg("pty") + isBackground := getBoolArg("background") + + if isPty { + if runtime.GOOS == "windows" { + return ErrorResult("PTY is not supported on Windows. Use background=true without pty.") + } + } + cwd := t.workingDir - if wd, ok := args["working_dir"].(string); ok && wd != "" { + if wd, ok := args["cwd"].(string); ok && wd != "" { if t.restrictToWorkspace && t.workingDir != "" { resolvedWD, err := validatePathWithAllowPaths(wd, t.workingDir, true, t.allowedPathPatterns) if err != nil { @@ -253,6 +344,14 @@ func (t *ExecTool) Execute(ctx context.Context, args map[string]any) *ToolResult } } + if isBackground { + return t.runBackground(ctx, command, cwd, isPty) + } + + return t.runSync(ctx, command, cwd) +} + +func (t *ExecTool) runSync(ctx context.Context, command, cwd string) *ToolResult { // timeout == 0 means no timeout var cmdCtx context.Context var cancel context.CancelFunc @@ -361,6 +460,560 @@ func (t *ExecTool) Execute(ctx context.Context, args map[string]any) *ToolResult } } +func (t *ExecTool) runBackground(ctx context.Context, command, cwd string, ptyEnabled bool) *ToolResult { + sessionID := generateSessionID() + session := &ProcessSession{ + ID: sessionID, + Command: command, + PTY: ptyEnabled, + Background: true, + StartTime: time.Now().Unix(), + Status: "running", + ptyKeyMode: PtyKeyModeCSI, + } + + var cmd *exec.Cmd + if runtime.GOOS == "windows" { + cmd = exec.Command("powershell", "-NoProfile", "-NonInteractive", "-Command", command) + } else { + cmd = exec.Command("sh", "-c", command) + } + if cwd != "" { + cmd.Dir = cwd + } + + prepareCommandForTermination(cmd) + + var stdoutReader io.ReadCloser + var stderrReader io.ReadCloser + var stdinWriter io.WriteCloser + + if ptyEnabled { + ptmx, tty, err := pty.Open() + if err != nil { + return ErrorResult(fmt.Sprintf("failed to create PTY: %v", err)) + } + + cmd.Stdin = tty + cmd.Stdout = tty + cmd.Stderr = tty + + // For PTY, we need Setsid to create a new session. + // Note: Setsid and Setpgid conflict, so we must replace SysProcAttr entirely. + setSysProcAttrForPty(cmd) + + session.ptyMaster = ptmx + } else { + var err error + stdoutReader, err = cmd.StdoutPipe() + if err != nil { + return ErrorResult(fmt.Sprintf("failed to create stdout pipe: %v", err)) + } + stderrReader, err = cmd.StderrPipe() + if err != nil { + return ErrorResult(fmt.Sprintf("failed to create stderr pipe: %v", err)) + } + stdinWriter, err = cmd.StdinPipe() + if err != nil { + return ErrorResult(fmt.Sprintf("failed to create stdin pipe: %v", err)) + } + session.stdoutPipe = io.MultiReader(stdoutReader, stderrReader) + session.stdinWriter = stdinWriter + } + + if err := cmd.Start(); err != nil { + if session.ptyMaster != nil { + session.ptyMaster.Close() + } + return ErrorResult(fmt.Sprintf("failed to start command: %v", err)) + } + + session.PID = cmd.Process.Pid + t.sessionManager.Add(session) + + session.outputBuffer = &bytes.Buffer{} + + // PTY mode: read from ptyMaster and wait for process + // Note: On Linux, closing ptyMaster doesn't interrupt blocking Read() calls, + // so we need cmd.Wait() in a separate goroutine to detect process exit. + if session.PTY && session.ptyMaster != nil { + go func() { + cmd.Wait() // Wait for process to exit + session.mu.Lock() + if cmd.ProcessState != nil { + session.ExitCode = cmd.ProcessState.ExitCode() + } + session.Status = "done" + session.mu.Unlock() + }() + + go func() { + buf := make([]byte, 4096) + for { + n, err := session.ptyMaster.Read(buf) + if n > 0 { + raw := string(buf[:n]) + if mode := detectPtyKeyMode(raw); mode != PtyKeyModeNotFound && mode != session.GetPtyKeyMode() { + session.SetPtyKeyMode(mode) + } + + session.mu.Lock() + if session.outputBuffer.Len() >= maxOutputBufferSize { + if !session.outputTruncated { + session.outputBuffer.WriteString(outputTruncateMarker) + session.outputTruncated = true + } + } else { + session.outputBuffer.Write(buf[:n]) + } + session.mu.Unlock() + } + if err != nil { + break + } + } + }() + } else { + // Non-PTY mode: single goroutine reads pipes. + // When Read() returns EOF (pipe closed), we break. + // When process exits, OS closes pipe write end → Read() returns EOF → we exit. + go func() { + buf := make([]byte, 4096) + + // Read stdout + for { + n, err := stdoutReader.Read(buf) + if n > 0 { + session.mu.Lock() + if session.outputBuffer.Len() >= maxOutputBufferSize { + if !session.outputTruncated { + session.outputBuffer.WriteString(outputTruncateMarker) + session.outputTruncated = true + } + } else { + session.outputBuffer.Write(buf[:n]) + } + session.mu.Unlock() + } + if err != nil { + break + } + } + + // Read stderr + for { + n, err := stderrReader.Read(buf) + if n > 0 { + session.mu.Lock() + if session.outputBuffer.Len() >= maxOutputBufferSize { + if !session.outputTruncated { + session.outputBuffer.WriteString(outputTruncateMarker) + session.outputTruncated = true + } + } else { + session.outputBuffer.Write(buf[:n]) + } + session.mu.Unlock() + } + if err != nil { + break + } + } + + // All pipes closed, get exit status + if stdinWriter != nil { + stdinWriter.Close() + } + cmd.Wait() + + session.mu.Lock() + if cmd.ProcessState != nil { + session.ExitCode = cmd.ProcessState.ExitCode() + } + session.Status = "done" + session.mu.Unlock() + }() + } + + resp := ExecResponse{ + SessionID: sessionID, + Status: "running", + } + data, _ := json.Marshal(resp) + return &ToolResult{ + ForLLM: string(data), + ForUser: fmt.Sprintf("Session %s started", sessionID), + IsError: false, + } +} + +func (t *ExecTool) executeList() *ToolResult { + sessions := t.sessionManager.List() + resp := ExecResponse{ + Sessions: sessions, + } + data, _ := json.Marshal(resp) + return &ToolResult{ + ForLLM: string(data), + ForUser: fmt.Sprintf("%d active sessions", len(sessions)), + IsError: false, + } +} + +func (t *ExecTool) executePoll(args map[string]any) *ToolResult { + sessionID, ok := args["sessionId"].(string) + if !ok { + return ErrorResult("sessionId is required") + } + + session, err := t.sessionManager.Get(sessionID) + if err != nil { + if errors.Is(err, ErrSessionNotFound) { + return ErrorResult(fmt.Sprintf("session not found: %s", sessionID)) + } + return ErrorResult(err.Error()) + } + + resp := ExecResponse{ + SessionID: sessionID, + Status: session.GetStatus(), + ExitCode: session.GetExitCode(), + } + data, _ := json.Marshal(resp) + return &ToolResult{ + ForLLM: string(data), + IsError: false, + } +} + +func (t *ExecTool) executeRead(args map[string]any) *ToolResult { + sessionID, ok := args["sessionId"].(string) + if !ok { + return ErrorResult("sessionId is required") + } + + session, err := t.sessionManager.Get(sessionID) + if err != nil { + if errors.Is(err, ErrSessionNotFound) { + return ErrorResult(fmt.Sprintf("session not found: %s", sessionID)) + } + return ErrorResult(err.Error()) + } + + output := session.Read() + + resp := ExecResponse{ + SessionID: sessionID, + Output: output, + Status: session.GetStatus(), + } + data, _ := json.Marshal(resp) + return &ToolResult{ + ForLLM: string(data), + IsError: false, + } +} + +func (t *ExecTool) executeWrite(args map[string]any) *ToolResult { + sessionID, ok := args["sessionId"].(string) + if !ok { + return ErrorResult("sessionId is required") + } + + data, ok := args["data"].(string) + if !ok { + return ErrorResult("data is required") + } + + session, err := t.sessionManager.Get(sessionID) + if err != nil { + if errors.Is(err, ErrSessionNotFound) { + return ErrorResult(fmt.Sprintf("session not found: %s", sessionID)) + } + return ErrorResult(err.Error()) + } + + if session.IsDone() { + return ErrorResult(fmt.Sprintf("process already exited with code %d", session.GetExitCode())) + } + + if err := session.Write(data); err != nil { + if errors.Is(err, ErrSessionDone) { + return ErrorResult(fmt.Sprintf("process already exited with code %d", session.GetExitCode())) + } + return ErrorResult(fmt.Sprintf("failed to write to session: %v", err)) + } + + resp := ExecResponse{ + SessionID: sessionID, + Status: session.GetStatus(), + } + respData, _ := json.Marshal(resp) + return &ToolResult{ + ForLLM: string(respData), + IsError: false, + } +} + +func (t *ExecTool) executeKill(args map[string]any) *ToolResult { + sessionID, ok := args["sessionId"].(string) + if !ok { + return ErrorResult("sessionId is required") + } + + session, err := t.sessionManager.Get(sessionID) + if err != nil { + if errors.Is(err, ErrSessionNotFound) { + return ErrorResult(fmt.Sprintf("session not found: %s", sessionID)) + } + return ErrorResult(err.Error()) + } + + if session.IsDone() { + return ErrorResult(fmt.Sprintf("process already exited with code %d", session.GetExitCode())) + } + + if err := session.Kill(); err != nil { + return ErrorResult(fmt.Sprintf("failed to kill session: %v", err)) + } + + t.sessionManager.Remove(sessionID) + + resp := ExecResponse{ + SessionID: sessionID, + Status: "done", + } + data, _ := json.Marshal(resp) + return &ToolResult{ + ForLLM: string(data), + ForUser: fmt.Sprintf("Session %s killed", sessionID), + IsError: false, + } +} + +// keyMap maps key names to their escape sequences. +var keyMap = map[string]string{ + "enter": "\r", + "return": "\r", + "tab": "\t", + "escape": "\x1b", + "esc": "\x1b", + "space": " ", + "backspace": "\x7f", + "bspace": "\x7f", + "up": "\x1b[A", + "down": "\x1b[B", + "right": "\x1b[C", + "left": "\x1b[D", + "home": "\x1b[1~", + "end": "\x1b[4~", + "pageup": "\x1b[5~", + "pagedown": "\x1b[6~", + "pgup": "\x1b[5~", + "pgdn": "\x1b[6~", + "insert": "\x1b[2~", + "ic": "\x1b[2~", + "delete": "\x1b[3~", + "del": "\x1b[3~", + "dc": "\x1b[3~", + "btab": "\x1b[Z", + "f1": "\x1bOP", + "f2": "\x1bOQ", + "f3": "\x1bOR", + "f4": "\x1bOS", + "f5": "\x1b[15~", + "f6": "\x1b[17~", + "f7": "\x1b[18~", + "f8": "\x1b[19~", + "f9": "\x1b[20~", + "f10": "\x1b[21~", + "f11": "\x1b[23~", + "f12": "\x1b[24~", +} + +// ss3KeysMap maps key names to SS3 escape sequences +var ss3KeysMap = map[string]string{ + "up": "\x1bOA", + "down": "\x1bOB", + "right": "\x1bOC", + "left": "\x1bOD", + "home": "\x1bOH", + "end": "\x1bOF", +} + +func detectPtyKeyMode(raw string) PtyKeyMode { + const SMKX = "\x1b[?1h" + const RMKX = "\x1b[?1l" + + lastSmkx := strings.LastIndex(raw, SMKX) + lastRmkx := strings.LastIndex(raw, RMKX) + + if lastSmkx == -1 && lastRmkx == -1 { + return PtyKeyModeNotFound + } + + if lastSmkx > lastRmkx { + return PtyKeyModeSS3 + } + return PtyKeyModeCSI +} + +// encodeKeyToken encodes a single key token into its escape sequence. +// Supports: +// - Named keys: "enter", "tab", "up", "ctrl-c", "alt-x", etc. +// - Ctrl modifier: "ctrl-c" or "c-c" (sends Ctrl+char) +// - Alt modifier: "alt-x" or "m-x" (sends ESC+char) +func encodeKeyToken(token string, ptyKeyMode PtyKeyMode) (string, error) { + token = strings.ToLower(strings.TrimSpace(token)) + if token == "" { + return "", nil + } + + // Handle ctrl-X format (c-x) + if strings.HasPrefix(token, "c-") { + char := token[2] + if char >= 'a' && char <= 'z' { + return string(rune(char) & 0x1f), nil // ctrl-a through ctrl-z + } + return "", fmt.Errorf("invalid ctrl key: %s", token) + } + + // Handle ctrl-X format (ctrl-x) + if strings.HasPrefix(token, "ctrl-") { + char := token[5] + if char >= 'a' && char <= 'z' { + return string(rune(char) & 0x1f), nil + } + return "", fmt.Errorf("invalid ctrl key: %s", token) + } + + // Handle alt-X format (m-x or alt-x) + if strings.HasPrefix(token, "m-") || strings.HasPrefix(token, "alt-") { + var char string + if strings.HasPrefix(token, "m-") { + char = token[2:] + } else { + char = token[4:] + } + if len(char) == 1 { + return "\x1b" + char, nil + } + return "", fmt.Errorf("invalid alt key: %s", token) + } + + // Handle shift modifier for special keys (shift-up, shift-down, etc.) + if strings.HasPrefix(token, "s-") || strings.HasPrefix(token, "shift-") { + var key string + if strings.HasPrefix(token, "s-") { + key = token[2:] + } else { + key = token[6:] + } + // Apply shift modifier: for single-char keys, return uppercase + if seq, ok := keyMap[key]; ok { + // For escape sequences, we can't easily add shift + // For single-char keys (letters), return uppercase + if len(seq) == 1 { + return strings.ToUpper(seq), nil + } + return seq, nil + } + return "", fmt.Errorf("unknown key with shift: %s", key) + } + + if ptyKeyMode == PtyKeyModeSS3 { + if seq, ok := ss3KeysMap[token]; ok { + return seq, nil + } + } + + if seq, ok := keyMap[token]; ok { + return seq, nil + } + + return "", fmt.Errorf("unknown key: %s (use write action for text input)", token) +} + +// encodeKeySequence encodes a slice of key tokens into a single string. +func encodeKeySequence(tokens []string, ptyKeyMode PtyKeyMode) (string, error) { + var result string + for _, token := range tokens { + seq, err := encodeKeyToken(token, ptyKeyMode) + if err != nil { + return "", err + } + result += seq + } + return result, nil +} + +func (t *ExecTool) executeSendKeys(args map[string]any) *ToolResult { + sessionID, ok := args["sessionId"].(string) + if !ok { + return ErrorResult("sessionId is required") + } + + keysStr, ok := args["keys"].(string) + if !ok { + return ErrorResult("keys must be a string") + } + + if keysStr == "" { + return ErrorResult("keys cannot be empty") + } + + // Parse comma-separated key names + keyNames := strings.Split(keysStr, ",") + var keys []string + for _, k := range keyNames { + k = strings.TrimSpace(k) + if k != "" { + keys = append(keys, k) + } + } + + if len(keys) == 0 { + return ErrorResult("keys cannot be empty") + } + + session, err := t.sessionManager.Get(sessionID) + if err != nil { + if errors.Is(err, ErrSessionNotFound) { + return ErrorResult(fmt.Sprintf("session not found: %s", sessionID)) + } + return ErrorResult(err.Error()) + } + + ptyKeyMode := session.GetPtyKeyMode() + + data, err := encodeKeySequence(keys, ptyKeyMode) + if err != nil { + return ErrorResult(fmt.Sprintf("invalid key: %v", err)) + } + + if session.IsDone() { + return ErrorResult(fmt.Sprintf("process already exited with code %d", session.GetExitCode())) + } + + if err := session.Write(data); err != nil { + if errors.Is(err, ErrSessionDone) { + return ErrorResult(fmt.Sprintf("process already exited with code %d", session.GetExitCode())) + } + return ErrorResult(fmt.Sprintf("failed to send keys: %v", err)) + } + + resp := ExecResponse{ + SessionID: sessionID, + Status: "running", + Output: fmt.Sprintf("Sent keys: %v", keys), + } + respData, _ := json.Marshal(resp) + return &ToolResult{ + ForLLM: string(respData), + IsError: false, + } +} + func (t *ExecTool) guardCommand(command, cwd string) string { cmd := strings.TrimSpace(command) lower := strings.ToLower(cmd) diff --git a/pkg/tools/shell_test.go b/pkg/tools/shell_test.go index f8f83ea74..a8de2f4c9 100644 --- a/pkg/tools/shell_test.go +++ b/pkg/tools/shell_test.go @@ -2,12 +2,16 @@ package tools import ( "context" + "encoding/json" "os" "path/filepath" + "runtime" "strings" "testing" "time" + "github.com/stretchr/testify/require" + "github.com/sipeed/picoclaw/pkg/config" ) @@ -20,6 +24,7 @@ func TestShellTool_Success(t *testing.T) { ctx := context.Background() args := map[string]any{ + "action": "run", "command": "echo 'hello world'", } @@ -50,6 +55,7 @@ func TestShellTool_Failure(t *testing.T) { ctx := context.Background() args := map[string]any{ + "action": "run", "command": "ls /nonexistent_directory_12345", } @@ -82,6 +88,7 @@ func TestShellTool_Timeout(t *testing.T) { ctx := context.Background() args := map[string]any{ + "action": "run", "command": "sleep 10", } @@ -112,8 +119,9 @@ func TestShellTool_WorkingDir(t *testing.T) { ctx := context.Background() args := map[string]any{ - "command": "cat test.txt", - "working_dir": tmpDir, + "action": "run", + "command": "cat test.txt", + "cwd": tmpDir, } result := tool.Execute(ctx, args) @@ -136,6 +144,7 @@ func TestShellTool_DangerousCommand(t *testing.T) { ctx := context.Background() args := map[string]any{ + "action": "run", "command": "rm -rf /", } @@ -159,6 +168,7 @@ func TestShellTool_DangerousCommand_KillBlocked(t *testing.T) { ctx := context.Background() args := map[string]any{ + "action": "run", "command": "kill 12345", } @@ -198,6 +208,7 @@ func TestShellTool_StderrCapture(t *testing.T) { ctx := context.Background() args := map[string]any{ + "action": "run", "command": "sh -c 'echo stdout; echo stderr >&2'", } @@ -222,6 +233,7 @@ func TestShellTool_OutputTruncation(t *testing.T) { ctx := context.Background() // Generate long output (>10000 chars) args := map[string]any{ + "action": "run", "command": "python3 -c \"print('x' * 20000)\" || echo " + strings.Repeat("x", 20000), } @@ -251,8 +263,9 @@ func TestShellTool_WorkingDir_OutsideWorkspace(t *testing.T) { } result := tool.Execute(context.Background(), map[string]any{ - "command": "pwd", - "working_dir": outsideDir, + "action": "run", + "command": "pwd", + "cwd": outsideDir, }) if !result.IsError { @@ -289,8 +302,9 @@ func TestShellTool_WorkingDir_SymlinkEscape(t *testing.T) { } result := tool.Execute(context.Background(), map[string]any{ - "command": "cat secret.txt", - "working_dir": link, + "action": "run", + "command": "cat secret.txt", + "cwd": link, }) if !result.IsError { @@ -312,7 +326,7 @@ func TestShellTool_RemoteChannelBlockedByDefault(t *testing.T) { t.Fatalf("NewExecToolWithConfig() error: %v", err) } ctx := WithToolContext(context.Background(), "telegram", "chat-1") - result := tool.Execute(ctx, map[string]any{"command": "echo hi"}) + result := tool.Execute(ctx, map[string]any{"action": "run", "command": "echo hi"}) if !result.IsError { t.Fatal("expected remote-channel exec to be blocked") @@ -333,7 +347,7 @@ func TestShellTool_InternalChannelAllowed(t *testing.T) { t.Fatalf("NewExecToolWithConfig() error: %v", err) } ctx := WithToolContext(context.Background(), "cli", "direct") - result := tool.Execute(ctx, map[string]any{"command": "echo hi"}) + result := tool.Execute(ctx, map[string]any{"action": "run", "command": "echo hi"}) if result.IsError { t.Fatalf("expected internal channel exec to succeed, got: %s", result.ForLLM) @@ -373,7 +387,7 @@ func TestShellTool_AllowRemoteBypassesChannelCheck(t *testing.T) { t.Fatalf("NewExecToolWithConfig() error: %v", err) } ctx := WithToolContext(context.Background(), "telegram", "chat-1") - result := tool.Execute(ctx, map[string]any{"command": "echo hi"}) + result := tool.Execute(ctx, map[string]any{"action": "run", "command": "echo hi"}) if result.IsError { t.Fatalf("expected allowRemote=true to permit remote channel, got: %s", result.ForLLM) @@ -392,6 +406,7 @@ func TestShellTool_RestrictToWorkspace(t *testing.T) { ctx := context.Background() args := map[string]any{ + "action": "run", "command": "cat ../../etc/passwd", } @@ -429,7 +444,7 @@ func TestShellTool_DevNullAllowed(t *testing.T) { } for _, cmd := range commands { - result := tool.Execute(context.Background(), map[string]any{"command": cmd}) + result := tool.Execute(context.Background(), map[string]any{"action": "run", "command": cmd}) if result.IsError && strings.Contains(result.ForLLM, "blocked") { t.Errorf("command should not be blocked: %s\n error: %s", cmd, result.ForLLM) } @@ -458,7 +473,7 @@ func TestShellTool_BlockDevices(t *testing.T) { } for _, cmd := range blocked { - result := tool.Execute(context.Background(), map[string]any{"command": cmd}) + result := tool.Execute(context.Background(), map[string]any{"action": "run", "command": cmd}) if !result.IsError { t.Errorf("expected block device write to be blocked: %s", cmd) } @@ -482,7 +497,7 @@ func TestShellTool_SafePathsInWorkspaceRestriction(t *testing.T) { } for _, cmd := range commands { - result := tool.Execute(context.Background(), map[string]any{"command": cmd}) + result := tool.Execute(context.Background(), map[string]any{"action": "run", "command": cmd}) if result.IsError && strings.Contains(result.ForLLM, "path outside working dir") { t.Errorf("safe path should not be blocked by workspace check: %s\n error: %s", cmd, result.ForLLM) } @@ -498,6 +513,7 @@ func TestShellTool_ExitCodeDetails(t *testing.T) { ctx := context.Background() args := map[string]any{ + "action": "run", "command": "sh -c 'exit 42'", } @@ -534,6 +550,7 @@ func TestShellTool_TimeoutWithPartialOutput(t *testing.T) { ctx := context.Background() // Use a command that outputs immediately then sleeps args := map[string]any{ + "action": "run", "command": "echo 'partial output before timeout' && sleep 30", } @@ -608,7 +625,9 @@ func TestShellTool_URLsNotBlocked(t *testing.T) { } for _, cmd := range commands { - result := tool.Execute(context.Background(), map[string]any{"command": cmd}) + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + result := tool.Execute(ctx, map[string]any{"action": "run", "command": cmd}) + cancel() if result.IsError && strings.Contains(result.ForLLM, "path outside working dir") { t.Errorf("command with URL should not be blocked by workspace check: %s\n error: %s", cmd, result.ForLLM) } @@ -633,7 +652,7 @@ func TestShellTool_FileURISandboxing(t *testing.T) { } for _, cmd := range blockedCommands { - result := tool.Execute(context.Background(), map[string]any{"command": cmd}) + result := tool.Execute(context.Background(), map[string]any{"action": "run", "command": cmd}) if !result.IsError || !strings.Contains(result.ForLLM, "path outside working dir") { t.Errorf("file:// URI outside workspace should be blocked: %s", cmd) } @@ -651,7 +670,7 @@ func TestShellTool_FileURISandboxing(t *testing.T) { } for _, cmd := range allowedCommands { - result := tool.Execute(context.Background(), map[string]any{"command": cmd}) + result := tool.Execute(context.Background(), map[string]any{"action": "run", "command": cmd}) if result.IsError && strings.Contains(result.ForLLM, "path outside working dir") { t.Errorf("file:// URI inside workspace should be allowed: %s\n error: %s", cmd, result.ForLLM) } @@ -677,9 +696,920 @@ func TestShellTool_URLBypassPrevented(t *testing.T) { } for _, cmd := range blockedCommands { - result := tool.Execute(context.Background(), map[string]any{"command": cmd}) + result := tool.Execute(context.Background(), map[string]any{"action": "run", "command": cmd}) if !result.IsError || !strings.Contains(result.ForLLM, "path outside working dir") { t.Errorf("bypass attempt should be blocked: %q\n got: %s", cmd, result.ForLLM) } } } + +func TestShellTool_Background_ReturnsImmediately(t *testing.T) { + tool, err := NewExecTool("", false) + require.NoError(t, err) + + ctx := context.Background() + args := map[string]any{ + "action": "run", + "command": "sleep 5", + "background": "true", + } + + start := time.Now() + result := tool.Execute(ctx, args) + elapsed := time.Since(start) + + require.False(t, result.IsError, "background run should not error: %s", result.ForLLM) + require.Less(t, elapsed, time.Second, "background run should return immediately") + require.Contains(t, result.ForLLM, "sessionId") +} + +func TestShellTool_List_Empty(t *testing.T) { + tool, err := NewExecTool("", false) + require.NoError(t, err) + + sm := NewSessionManager() + tool.sessionManager = sm + + ctx := context.Background() + args := map[string]any{"action": "list"} + + result := tool.Execute(ctx, args) + require.False(t, result.IsError) + require.Contains(t, result.ForUser, "0 active sessions") +} + +func TestShellTool_RunBackground_List(t *testing.T) { + tool, err := NewExecTool("", false) + require.NoError(t, err) + + sm := NewSessionManager() + tool.sessionManager = sm + + ctx := WithToolContext(context.Background(), "cli", "test") + + runResult := tool.Execute(ctx, map[string]any{ + "action": "run", + "command": "sleep 10", + "background": "true", + }) + require.False(t, runResult.IsError, "run should succeed: %s", runResult.ForLLM) + + var resp ExecResponse + err = json.Unmarshal([]byte(runResult.ForLLM), &resp) + require.NoError(t, err) + require.NotEmpty(t, resp.SessionID) + + time.Sleep(100 * time.Millisecond) + + listResult := tool.Execute(ctx, map[string]any{"action": "list"}) + require.False(t, listResult.IsError) + + var listResp ExecResponse + err = json.Unmarshal([]byte(listResult.ForLLM), &listResp) + require.NoError(t, err) + require.Len(t, listResp.Sessions, 1) + require.Equal(t, resp.SessionID, listResp.Sessions[0].ID) + + killResult := tool.Execute(ctx, map[string]any{ + "action": "kill", + "sessionId": resp.SessionID, + }) + require.False(t, killResult.IsError, "kill should succeed: %s", killResult.ForLLM) +} + +func TestShellTool_Read_Output(t *testing.T) { + tool, err := NewExecTool("", false) + require.NoError(t, err) + + sm := NewSessionManager() + tool.sessionManager = sm + + ctx := WithToolContext(context.Background(), "cli", "test") + + runResult := tool.Execute(ctx, map[string]any{ + "action": "run", + "command": "echo hello", + "background": "true", + }) + require.False(t, runResult.IsError) + + var resp ExecResponse + err = json.Unmarshal([]byte(runResult.ForLLM), &resp) + require.NoError(t, err) + + time.Sleep(200 * time.Millisecond) + + readResult := tool.Execute(ctx, map[string]any{ + "action": "read", + "sessionId": resp.SessionID, + }) + + if !readResult.IsError { + var readResp ExecResponse + err = json.Unmarshal([]byte(readResult.ForLLM), &readResp) + require.NoError(t, err) + } +} + +func TestShellTool_Kill(t *testing.T) { + tool, err := NewExecTool("", false) + require.NoError(t, err) + + sm := NewSessionManager() + tool.sessionManager = sm + + ctx := WithToolContext(context.Background(), "cli", "test") + + runResult := tool.Execute(ctx, map[string]any{ + "action": "run", + "command": "sleep 100", + "background": "true", + }) + require.False(t, runResult.IsError) + + var resp ExecResponse + err = json.Unmarshal([]byte(runResult.ForLLM), &resp) + require.NoError(t, err) + + killResult := tool.Execute(ctx, map[string]any{ + "action": "kill", + "sessionId": resp.SessionID, + }) + require.False(t, killResult.IsError, "kill should succeed: %s", killResult.ForLLM) + + time.Sleep(100 * time.Millisecond) + + listResult := tool.Execute(ctx, map[string]any{"action": "list"}) + var listResp ExecResponse + err = json.Unmarshal([]byte(listResult.ForLLM), &listResp) + require.NoError(t, err) + require.Len(t, listResp.Sessions, 0) +} + +func TestShellTool_PTY_AllowedCommands(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("PTY not supported on Windows") + } + + tool, err := NewExecTool("", false) + require.NoError(t, err) + + sm := NewSessionManager() + tool.sessionManager = sm + + ctx := WithToolContext(context.Background(), "cli", "test") + + // Test that PTY is allowed for non-interpreter commands + result := tool.Execute(ctx, map[string]any{ + "action": "run", + "command": "cat", + "pty": "true", + "background": "true", + }) + require.False(t, result.IsError, "PTY with cat should succeed: %s", result.ForLLM) + require.Contains(t, result.ForLLM, "sessionId") + + var resp ExecResponse + err = json.Unmarshal([]byte(result.ForLLM), &resp) + require.NoError(t, err) + require.NotEmpty(t, resp.SessionID) + + // Clean up + tool.Execute(ctx, map[string]any{ + "action": "kill", + "sessionId": resp.SessionID, + }) +} + +func TestShellTool_PTY_WriteRead(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("PTY not supported on Windows") + } + + tool, err := NewExecTool("", false) + require.NoError(t, err) + + sm := NewSessionManager() + tool.sessionManager = sm + + ctx := WithToolContext(context.Background(), "cli", "test") + + // Start a PTY session with a command that waits for input + // Using 'cat' which will wait for stdin + result := tool.Execute(ctx, map[string]any{ + "action": "run", + "command": "cat", + "pty": "true", + "background": "true", + }) + require.False(t, result.IsError, "PTY run should succeed: %s", result.ForLLM) + + var resp ExecResponse + err = json.Unmarshal([]byte(result.ForLLM), &resp) + require.NoError(t, err) + + // Write some input to cat + writeResult := tool.Execute(ctx, map[string]any{ + "action": "write", + "sessionId": resp.SessionID, + "data": "hello\n", + }) + require.False(t, writeResult.IsError, "write should succeed: %s", writeResult.ForLLM) + + // Give cat time to process and output + time.Sleep(200 * time.Millisecond) + + // Read the output + readResult := tool.Execute(ctx, map[string]any{ + "action": "read", + "sessionId": resp.SessionID, + }) + + require.False(t, readResult.IsError, "read should succeed: %s", readResult.ForLLM) + + var readResp ExecResponse + err = json.Unmarshal([]byte(readResult.ForLLM), &readResp) + require.NoError(t, err) + // PTY output should contain "hello" + require.Contains(t, readResp.Output, "hello") + + // Clean up + tool.Execute(ctx, map[string]any{ + "action": "kill", + "sessionId": resp.SessionID, + }) +} + +func TestShellTool_PTY_Poll(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("PTY not supported on Windows") + } + + tool, err := NewExecTool("", false) + require.NoError(t, err) + + sm := NewSessionManager() + tool.sessionManager = sm + + ctx := WithToolContext(context.Background(), "cli", "test") + + // Start a PTY session with a long-running command + result := tool.Execute(ctx, map[string]any{ + "action": "run", + "command": "sleep 2", + "pty": "true", + "background": "true", + }) + require.False(t, result.IsError, "PTY run should succeed: %s", result.ForLLM) + + var resp ExecResponse + err = json.Unmarshal([]byte(result.ForLLM), &resp) + require.NoError(t, err) + + // Poll should show running + pollResult := tool.Execute(ctx, map[string]any{ + "action": "poll", + "sessionId": resp.SessionID, + }) + require.False(t, pollResult.IsError, "poll should succeed: %s", pollResult.ForLLM) + + var pollResp ExecResponse + err = json.Unmarshal([]byte(pollResult.ForLLM), &pollResp) + require.NoError(t, err) + require.Equal(t, "running", pollResp.Status) + + // Wait for sleep to complete + time.Sleep(2500 * time.Millisecond) + + // Poll should show done + pollResult = tool.Execute(ctx, map[string]any{ + "action": "poll", + "sessionId": resp.SessionID, + }) + require.False(t, pollResult.IsError) + + err = json.Unmarshal([]byte(pollResult.ForLLM), &pollResp) + require.NoError(t, err) + require.Equal(t, "done", pollResp.Status) +} + +func TestShellTool_PTY_Kill(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("PTY not supported on Windows") + } + + tool, err := NewExecTool("", false) + require.NoError(t, err) + + sm := NewSessionManager() + tool.sessionManager = sm + + ctx := WithToolContext(context.Background(), "cli", "test") + + // Start a PTY session with a long-running command + result := tool.Execute(ctx, map[string]any{ + "action": "run", + "command": "sleep 10", + "pty": "true", + "background": "true", + }) + require.False(t, result.IsError, "PTY run should succeed: %s", result.ForLLM) + + var resp ExecResponse + err = json.Unmarshal([]byte(result.ForLLM), &resp) + require.NoError(t, err) + + // Kill the session + killResult := tool.Execute(ctx, map[string]any{ + "action": "kill", + "sessionId": resp.SessionID, + }) + require.False(t, killResult.IsError, "kill should succeed: %s", killResult.ForLLM) + + // Verify kill response shows done status + var killResp ExecResponse + err = json.Unmarshal([]byte(killResult.ForLLM), &killResp) + require.NoError(t, err) + require.Equal(t, "done", killResp.Status) + + // Poll should return error since session is removed after kill + pollResult := tool.Execute(ctx, map[string]any{ + "action": "poll", + "sessionId": resp.SessionID, + }) + // Session is removed after kill, so poll returns error with "session not found" + require.True(t, pollResult.IsError, "poll should error after kill (session removed)") + require.Contains(t, pollResult.ForLLM, "session not found") +} + +func TestShellTool_Write_Read_NonPTY(t *testing.T) { + tool, err := NewExecTool("", false) + require.NoError(t, err) + + sm := NewSessionManager() + tool.sessionManager = sm + + ctx := WithToolContext(context.Background(), "cli", "test") + + // Start a background process that reads from stdin and outputs it + // Using 'cat' which echoes stdin to stdout + result := tool.Execute(ctx, map[string]any{ + "action": "run", + "command": "cat", + "pty": false, + "background": "true", + }) + require.False(t, result.IsError, "run should succeed: %s", result.ForLLM) + + var resp ExecResponse + err = json.Unmarshal([]byte(result.ForLLM), &resp) + require.NoError(t, err) + + // Write some input to cat + writeResult := tool.Execute(ctx, map[string]any{ + "action": "write", + "sessionId": resp.SessionID, + "data": "hello world\n", + }) + require.False(t, writeResult.IsError, "write should succeed: %s", writeResult.ForLLM) + + // Give cat time to process and output + time.Sleep(200 * time.Millisecond) + + // Read the output + readResult := tool.Execute(ctx, map[string]any{ + "action": "read", + "sessionId": resp.SessionID, + }) + require.False(t, readResult.IsError, "read should succeed: %s", readResult.ForLLM) + + var readResp ExecResponse + err = json.Unmarshal([]byte(readResult.ForLLM), &readResp) + require.NoError(t, err) + require.Contains(t, readResp.Output, "hello world") + + // Clean up + tool.Execute(ctx, map[string]any{ + "action": "kill", + "sessionId": resp.SessionID, + }) +} + +func TestShellTool_Read_NonPTY_Running(t *testing.T) { + tool, err := NewExecTool("", false) + require.NoError(t, err) + + sm := NewSessionManager() + tool.sessionManager = sm + + ctx := WithToolContext(context.Background(), "cli", "test") + + // Start a long-running process that produces output over time + // Using sh -c with sleep at the end so process doesn't exit immediately + result := tool.Execute(ctx, map[string]any{ + "action": "run", + "command": "sh -c 'echo line1; sleep 0.5; echo line2; sleep 0.5; echo line3; sleep 10'", + "pty": false, + "background": "true", + }) + require.False(t, result.IsError, "run should succeed: %s", result.ForLLM) + + var resp ExecResponse + err = json.Unmarshal([]byte(result.ForLLM), &resp) + require.NoError(t, err) + + // Give time for first outputs to be produced + time.Sleep(300 * time.Millisecond) + + // Read output while process is running + readResult := tool.Execute(ctx, map[string]any{ + "action": "read", + "sessionId": resp.SessionID, + }) + require.False(t, readResult.IsError, "read should succeed: %s", readResult.ForLLM) + + var readResp ExecResponse + err = json.Unmarshal([]byte(readResult.ForLLM), &readResp) + require.NoError(t, err) + // Should have at least line1 + require.Contains(t, readResp.Output, "line1") + + // Wait for line3 to be produced (line1=0s, line2=0.5s, line3=1s, then sleep 10) + time.Sleep(1200 * time.Millisecond) + + // Read again - should have line3 as well + readResult = tool.Execute(ctx, map[string]any{ + "action": "read", + "sessionId": resp.SessionID, + }) + require.False(t, readResult.IsError, "read should succeed: %s", readResult.ForLLM) + + err = json.Unmarshal([]byte(readResult.ForLLM), &readResp) + require.NoError(t, err) + require.Contains(t, readResp.Output, "line3") + + // Clean up + tool.Execute(ctx, map[string]any{ + "action": "kill", + "sessionId": resp.SessionID, + }) +} + +func TestShellTool_ProcessGroupKill(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("Process group kill not supported on Windows") + } + + // Note: Testing process group kill with PTY is tricky because the command + // must be run through an interpreter (sh, bash) which is blocked for PTY. + // Instead, we test with non-PTY mode which also uses Setsid for background processes. + + tool, err := NewExecTool("", false) + require.NoError(t, err) + + sm := NewSessionManager() + tool.sessionManager = sm + + ctx := WithToolContext(context.Background(), "cli", "test") + + // Start a shell that spawns child processes (non-PTY mode) + // The sh -c command creates child sleep processes + result := tool.Execute(ctx, map[string]any{ + "action": "run", + "command": "sh -c 'sleep 30 & sleep 30 & wait'", + "pty": false, + "background": "true", + }) + require.False(t, result.IsError, "run should succeed: %s", result.ForLLM) + + var resp ExecResponse + err = json.Unmarshal([]byte(result.ForLLM), &resp) + require.NoError(t, err) + + // Give time for child processes to spawn + time.Sleep(500 * time.Millisecond) + + // Kill the session - should kill the entire process group + killResult := tool.Execute(ctx, map[string]any{ + "action": "kill", + "sessionId": resp.SessionID, + }) + require.False(t, killResult.IsError, "kill should succeed: %s", killResult.ForLLM) + + // Verify kill response shows done status + var killResp ExecResponse + err = json.Unmarshal([]byte(killResult.ForLLM), &killResp) + require.NoError(t, err) + require.Equal(t, "done", killResp.Status) + + // Poll should return error since session is removed after kill + pollResult := tool.Execute(ctx, map[string]any{ + "action": "poll", + "sessionId": resp.SessionID, + }) + require.True(t, pollResult.IsError, "poll should error after kill (session removed)") + require.Contains(t, pollResult.ForLLM, "session not found") +} + +func TestShellTool_PTY_ProcessGroupKill(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("PTY process group kill not supported on Windows") + } + + // This test binary creates 4 child sleep processes and waits for signals. + // It's not an interpreter, so it's allowed with PTY mode. + // The binary is created in /tmp/test_pgroup.c and compiled as part of test setup. + testBinary := "/tmp/test_pgroup" + if _, err := os.Stat(testBinary); os.IsNotExist(err) { + t.Skip("Test binary /tmp/test_pgroup not found - run: gcc -o /tmp/test_pgroup /tmp/test_pgroup.c") + } + + tool, err := NewExecTool("", false) + require.NoError(t, err) + + sm := NewSessionManager() + tool.sessionManager = sm + + ctx := WithToolContext(context.Background(), "cli", "test") + + // Start the test binary with PTY mode + // It forks 4 child sleep processes and waits for signals + result := tool.Execute(ctx, map[string]any{ + "action": "run", + "command": testBinary, + "pty": "true", + "background": "true", + }) + require.False(t, result.IsError, "run should succeed: %s", result.ForLLM) + + var resp ExecResponse + err = json.Unmarshal([]byte(result.ForLLM), &resp) + require.NoError(t, err) + + // Give time for child processes to spawn + time.Sleep(500 * time.Millisecond) + + // Kill the session - should kill the entire process group + killResult := tool.Execute(ctx, map[string]any{ + "action": "kill", + "sessionId": resp.SessionID, + }) + require.False(t, killResult.IsError, "kill should succeed: %s", killResult.ForLLM) + + // Verify kill response shows done status + var killResp ExecResponse + err = json.Unmarshal([]byte(killResult.ForLLM), &killResp) + require.NoError(t, err) + require.Equal(t, "done", killResp.Status) + + // Poll should return error since session is removed after kill + pollResult := tool.Execute(ctx, map[string]any{ + "action": "poll", + "sessionId": resp.SessionID, + }) + require.True(t, pollResult.IsError, "poll should error after kill (session removed)") + require.Contains(t, pollResult.ForLLM, "session not found") +} + +func TestShellTool_PTY_Background_Read(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("PTY not supported on Windows") + } + + tool, err := NewExecTool("", false) + require.NoError(t, err) + + sm := NewSessionManager() + tool.sessionManager = sm + + ctx := WithToolContext(context.Background(), "cli", "test") + + // Start a fast command with PTY + background mode + runResult := tool.Execute(ctx, map[string]any{ + "action": "run", + "command": "echo hello", + "pty": "true", + "background": "true", + }) + require.False(t, runResult.IsError, "run should succeed: %s", runResult.ForLLM) + + var runResp ExecResponse + err = json.Unmarshal([]byte(runResult.ForLLM), &runResp) + require.NoError(t, err) + require.NotEmpty(t, runResp.SessionID) + require.Equal(t, "running", runResp.Status) + + // Wait for command to complete + time.Sleep(500 * time.Millisecond) + + // Read output - this is the key test: PTY + background mode should preserve output + readResult := tool.Execute(ctx, map[string]any{ + "action": "read", + "sessionId": runResp.SessionID, + }) + require.False(t, readResult.IsError, "read should succeed: %s", readResult.ForLLM) + require.Contains(t, readResult.ForLLM, "hello", "output should contain 'hello'") +} + +func TestShellTool_PTY_Background_ReadNoBlock(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("PTY not supported on Windows") + } + + tool, err := NewExecTool("", false) + require.NoError(t, err) + + sm := NewSessionManager() + tool.sessionManager = sm + + ctx := WithToolContext(context.Background(), "cli", "test") + + // Start a long-running command with PTY + background mode + // This command produces no output, just sleeps + runResult := tool.Execute(ctx, map[string]any{ + "action": "run", + "command": "sleep 10", + "pty": "true", + "background": "true", + }) + require.False(t, runResult.IsError, "run should succeed: %s", runResult.ForLLM) + + var runResp ExecResponse + err = json.Unmarshal([]byte(runResult.ForLLM), &runResp) + require.NoError(t, err) + require.NotEmpty(t, runResp.SessionID) + + // Read immediately - should NOT block even though process is running and has no output + // This tests that Read() returns quickly (within 1 second) instead of blocking for 10 seconds + start := time.Now() + readResult := tool.Execute(ctx, map[string]any{ + "action": "read", + "sessionId": runResp.SessionID, + }) + elapsed := time.Since(start) + + require.False(t, readResult.IsError, "read should succeed: %s", readResult.ForLLM) + require.Less(t, elapsed.Seconds(), 1.0, "read should not block, should return within 1 second") + + // Kill the session to clean up + killResult := tool.Execute(ctx, map[string]any{ + "action": "kill", + "sessionId": runResp.SessionID, + }) + require.False(t, killResult.IsError, "kill should succeed: %s", killResult.ForLLM) +} + +func TestShellTool_Poll_Status(t *testing.T) { + tool, err := NewExecTool("", false) + require.NoError(t, err) + + sm := NewSessionManager() + tool.sessionManager = sm + + ctx := WithToolContext(context.Background(), "cli", "test") + + runResult := tool.Execute(ctx, map[string]any{ + "action": "run", + "command": "sleep 1", + "background": "true", + }) + require.False(t, runResult.IsError) + + var resp ExecResponse + err = json.Unmarshal([]byte(runResult.ForLLM), &resp) + require.NoError(t, err) + + pollResult := tool.Execute(ctx, map[string]any{ + "action": "poll", + "sessionId": resp.SessionID, + }) + require.False(t, pollResult.IsError) + + var pollResp ExecResponse + err = json.Unmarshal([]byte(pollResult.ForLLM), &pollResp) + require.NoError(t, err) + require.Equal(t, "running", pollResp.Status) + + time.Sleep(1200 * time.Millisecond) + + pollResult = tool.Execute(ctx, map[string]any{ + "action": "poll", + "sessionId": resp.SessionID, + }) + require.False(t, pollResult.IsError) + + err = json.Unmarshal([]byte(pollResult.ForLLM), &pollResp) + require.NoError(t, err) + require.Equal(t, "done", pollResp.Status) +} + +func TestShellTool_Action_Run_Sync(t *testing.T) { + tool, err := NewExecTool("", false) + require.NoError(t, err) + + ctx := context.Background() + + result := tool.Execute(ctx, map[string]any{ + "action": "run", + "command": "echo hello", + }) + + require.False(t, result.IsError) + require.Contains(t, result.ForLLM, "hello") +} + +// TestShellTool_Background_ReadAfterExit verifies that we can read +// buffered output even after the background process has exited. +func TestShellTool_Background_ReadAfterExit(t *testing.T) { + tool, err := NewExecTool("", false) + require.NoError(t, err) + + ctx := context.Background() + + // Start a background command that produces output and exits quickly + runResult := tool.Execute(ctx, map[string]any{ + "action": "run", + "command": "echo hello && sleep 1 && echo world", + "background": "true", + }) + require.False(t, runResult.IsError, "run should succeed: %s", runResult.ForUser) + + // Parse session ID from response + var resp ExecResponse + err = json.Unmarshal([]byte(runResult.ForLLM), &resp) + require.NoError(t, err) + require.NotEmpty(t, resp.SessionID) + sessionID := resp.SessionID + + // Wait for process to exit (sleep 1 + some buffer) + time.Sleep(1500 * time.Millisecond) + + // Poll to verify process is done + pollResult := tool.Execute(ctx, map[string]any{ + "action": "poll", + "sessionId": sessionID, + }) + require.False(t, pollResult.IsError, "poll should succeed: %s", pollResult.ForLLM) + var pollResp ExecResponse + err = json.Unmarshal([]byte(pollResult.ForLLM), &pollResp) + require.NoError(t, err) + require.Equal(t, "done", pollResp.Status, "process should be done") + + // Try to read output AFTER process has exited + readResult := tool.Execute(ctx, map[string]any{ + "action": "read", + "sessionId": sessionID, + }) + require.False(t, readResult.IsError, "read should succeed after exit: %s", readResult.ForLLM) + + var readResp ExecResponse + err = json.Unmarshal([]byte(readResult.ForLLM), &readResp) + require.NoError(t, err) + + // Output should contain both "hello" and "world" + require.Contains(t, readResp.Output, "hello", "should contain hello") + require.Contains(t, readResp.Output, "world", "should contain world after sleep") +} + +func TestSendKeys_CtrlC(t *testing.T) { + // Note: Ctrl-C as a signal requires sending SIGINT to the process group, + // which requires elevated privileges. Writing "\x03" to PTY passes the byte + // to the process but doesn't generate SIGINT for processes that don't read stdin. + // For interrupting processes, use the kill action instead. + t.Skip("Ctrl-C as signal not supported - use kill action for interruption") +} + +func TestEncodeKeyToken(t *testing.T) { + tests := []struct { + token string + expected string + hasError bool + }{ + // Named keys + {"enter", "\r", false}, + {"return", "\r", false}, + {"tab", "\t", false}, + {"escape", "\x1b", false}, + {"esc", "\x1b", false}, + {"backspace", "\x7f", false}, + {"up", "\x1b[A", false}, + {"down", "\x1b[B", false}, + {"left", "\x1b[D", false}, + {"right", "\x1b[C", false}, + {"home", "\x1b[1~", false}, + {"end", "\x1b[4~", false}, + {"pageup", "\x1b[5~", false}, + {"pagedown", "\x1b[6~", false}, + {"delete", "\x1b[3~", false}, + {"f1", "\x1bOP", false}, + {"f12", "\x1b[24~", false}, + + // Ctrl keys + {"ctrl-c", "\x03", false}, + {"ctrl-d", "\x04", false}, + {"ctrl-a", "\x01", false}, + {"ctrl-z", "\x1a", false}, + {"c-c", "\x03", false}, + {"c-d", "\x04", false}, + + // Alt keys + {"alt-x", "\x1bx", false}, + {"m-x", "\x1bx", false}, + + // Case insensitive tests + {"ENTER", "\r", false}, + {"TAB", "\t", false}, + {"CTRL-C", "\x03", false}, + {"Ctrl-D", "\x04", false}, + {"ALT-X", "\x1bx", false}, + {"M-X", "\x1bx", false}, + {"UP", "\x1b[A", false}, + {"DOWN", "\x1b[B", false}, + + // Unknown keys should return error (use write action for text input) + {"unknown-key", "", true}, + } + + for _, tt := range tests { + t.Run(tt.token, func(t *testing.T) { + result, err := encodeKeyToken(tt.token, PtyKeyModeCSI) + if tt.hasError { + require.Error(t, err, "expected error for %s", tt.token) + } else { + require.NoError(t, err, "unexpected error for %s", tt.token) + require.Equal(t, tt.expected, result, "wrong encoding for %s", tt.token) + } + }) + } +} + +// TestDetectPtyKeyMode tests smkx/rmkx detection in PTY output +func TestDetectPtyKeyMode(t *testing.T) { + tests := []struct { + name string + raw string + expected PtyKeyMode + }{ + {"no toggle", "hello world", PtyKeyModeNotFound}, + {"smkx only", "\x1b[?1h\x1b=", PtyKeyModeSS3}, + {"rmkx only", "\x1b[?1l\x1b>", PtyKeyModeCSI}, + {"both smkx first", "\x1b[?1h\x1b=...\x1b[?1l\x1b>", PtyKeyModeCSI}, + {"both rmkx first", "\x1b[?1l\x1b>...\x1b[?1h\x1b=", PtyKeyModeSS3}, + {"multiple toggles smkx last", "\x1b[?1h\x1b=...\x1b[?1l\x1b>...\x1b[?1h\x1b=", PtyKeyModeSS3}, + {"multiple toggles rmkx last", "\x1b[?1l\x1b>...\x1b[?1h\x1b=...\x1b[?1l\x1b>", PtyKeyModeCSI}, + {"partial smkx", "\x1b[?1h", PtyKeyModeSS3}, + {"partial rmkx", "\x1b[?1l", PtyKeyModeCSI}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := detectPtyKeyMode(tt.raw) + require.Equal(t, tt.expected, result, "wrong mode for %s", tt.name) + }) + } +} + +func TestEncodeKeyTokenWithPtyKeyMode(t *testing.T) { + tests := []struct { + name string + token string + mode PtyKeyMode + expected string + hasError bool + }{ + // CSI mode + {"up csi", "up", PtyKeyModeCSI, "\x1b[A", false}, + {"down csi", "down", PtyKeyModeCSI, "\x1b[B", false}, + {"left csi", "left", PtyKeyModeCSI, "\x1b[D", false}, + {"right csi", "right", PtyKeyModeCSI, "\x1b[C", false}, + + // SS3 mode + {"up ss3", "up", PtyKeyModeSS3, "\x1bOA", false}, + {"down ss3", "down", PtyKeyModeSS3, "\x1bOB", false}, + {"left ss3", "left", PtyKeyModeSS3, "\x1bOD", false}, + {"right ss3", "right", PtyKeyModeSS3, "\x1bOC", false}, + {"home ss3", "home", PtyKeyModeSS3, "\x1bOH", false}, + {"end ss3", "end", PtyKeyModeSS3, "\x1bOF", false}, + + // Other keys unaffected by mode + {"enter ss3", "enter", PtyKeyModeSS3, "\r", false}, + {"tab ss3", "tab", PtyKeyModeSS3, "\t", false}, + {"ctrl-c ss3", "ctrl-c", PtyKeyModeSS3, "\x03", false}, + + // NotFound behaves like CSI + {"up notfound", "up", PtyKeyModeNotFound, "\x1b[A", false}, + {"down notfound", "down", PtyKeyModeNotFound, "\x1b[B", false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result, err := encodeKeyToken(tt.token, tt.mode) + if tt.hasError { + require.Error(t, err, "expected error for %s", tt.name) + } else { + require.NoError(t, err, "unexpected error for %s", tt.name) + require.Equal(t, tt.expected, result, "wrong encoding for %s", tt.name) + } + }) + } +} diff --git a/pkg/tools/shell_timeout_unix_test.go b/pkg/tools/shell_timeout_unix_test.go index 357e1276e..dfd28454c 100644 --- a/pkg/tools/shell_timeout_unix_test.go +++ b/pkg/tools/shell_timeout_unix_test.go @@ -30,6 +30,7 @@ func TestShellTool_TimeoutKillsChildProcess(t *testing.T) { tool.SetTimeout(500 * time.Millisecond) args := map[string]any{ + "action": "run", // Spawn a child process that would outlive the shell unless process-group kill is used. "command": "sleep 60 & echo $! > child.pid; wait", } diff --git a/pkg/tools/sysproc_unix.go b/pkg/tools/sysproc_unix.go new file mode 100644 index 000000000..0fb03d43a --- /dev/null +++ b/pkg/tools/sysproc_unix.go @@ -0,0 +1,12 @@ +//go:build !windows + +package tools + +import ( + "os/exec" + "syscall" +) + +func setSysProcAttrForPty(cmd *exec.Cmd) { + cmd.SysProcAttr = &syscall.SysProcAttr{Setsid: true} +} diff --git a/pkg/tools/sysproc_windows.go b/pkg/tools/sysproc_windows.go new file mode 100644 index 000000000..150f166fb --- /dev/null +++ b/pkg/tools/sysproc_windows.go @@ -0,0 +1,10 @@ +//go:build windows + +package tools + +import "os/exec" + +func setSysProcAttrForPty(cmd *exec.Cmd) { + // Windows doesn't support Setsid, and PTY is not available on Windows anyway. + // This function is a no-op for Windows builds. +} diff --git a/pkg/tools/toolloop.go b/pkg/tools/toolloop.go index 244f0d4a2..387813e94 100644 --- a/pkg/tools/toolloop.go +++ b/pkg/tools/toolloop.go @@ -159,10 +159,7 @@ func RunToolLoop( // Append results in original order for _, r := range results { - contentForLLM := r.result.ForLLM - if contentForLLM == "" && r.result.Err != nil { - contentForLLM = r.result.Err.Error() - } + contentForLLM := r.result.ContentForLLM() messages = append(messages, providers.Message{ Role: "tool", diff --git a/pkg/tools/types.go b/pkg/tools/types.go index a6015cde3..4d1a18d5a 100644 --- a/pkg/tools/types.go +++ b/pkg/tools/types.go @@ -56,3 +56,24 @@ type ToolFunctionDefinition struct { Description string `json:"description"` Parameters map[string]any `json:"parameters"` } + +type ExecRequest struct { + Action string `json:"action"` + Command string `json:"command,omitempty"` + PTY bool `json:"pty,omitempty"` + Background bool `json:"background,omitempty"` + Timeout int `json:"timeout,omitempty"` + Env map[string]string `json:"env,omitempty"` + Cwd string `json:"cwd,omitempty"` + SessionID string `json:"sessionId,omitempty"` + Data string `json:"data,omitempty"` +} + +type ExecResponse struct { + SessionID string `json:"sessionId,omitempty"` + Status string `json:"status,omitempty"` + ExitCode int `json:"exitCode,omitempty"` + Output string `json:"output,omitempty"` + Error string `json:"error,omitempty"` + Sessions []SessionInfo `json:"sessions,omitempty"` +} diff --git a/pkg/tools/web.go b/pkg/tools/web.go index 7ff724802..342f7458b 100644 --- a/pkg/tools/web.go +++ b/pkg/tools/web.go @@ -43,7 +43,9 @@ var ( reBlankLines = regexp.MustCompile(`\n{3,}`) // DuckDuckGo result extraction - reDDGLink = regexp.MustCompile(`]*class="[^"]*result__a[^"]*"[^>]*href="([^"]+)"[^>]*>([\s\S]*?)`) + reDDGLink = regexp.MustCompile( + `]*class="[^"]*result__a[^"]*"[^>]*href="([^"]+)"[^>]*>([\s\S]*?)`, + ) reDDGSnippet = regexp.MustCompile(`([\s\S]*?)`) ) @@ -86,7 +88,122 @@ func (it *APIKeyIterator) Next() (string, bool) { } type SearchProvider interface { - Search(ctx context.Context, query string, count int) (string, error) + Search(ctx context.Context, query string, count int, rangeCode string) (string, error) +} + +func normalizeSearchRange(raw string) (string, error) { + rangeCode := strings.ToLower(strings.TrimSpace(raw)) + switch rangeCode { + case "", "d", "w", "m", "y": + return rangeCode, nil + default: + return "", fmt.Errorf("range must be one of: d, w, m, y") + } +} + +func mapBraveFreshness(rangeCode string) string { + switch rangeCode { + case "d": + return "pd" + case "w": + return "pw" + case "m": + return "pm" + case "y": + return "py" + default: + return "" + } +} + +func mapTavilyTimeRange(rangeCode string) string { + switch rangeCode { + case "d": + return "day" + case "w": + return "week" + case "m": + return "month" + case "y": + return "year" + default: + return "" + } +} + +func mapPerplexityRecencyFilter(rangeCode string) string { + switch rangeCode { + case "d": + return "day" + case "w": + return "week" + case "m": + return "month" + case "y": + return "year" + default: + return "" + } +} + +func mapDuckDuckGoDateFilter(rangeCode string) string { + switch rangeCode { + case "d": + return "d" + case "w": + return "w" + case "m": + return "m" + case "y": + return "t" + default: + return "" + } +} + +func mapSearXNGTimeRange(rangeCode string) string { + switch rangeCode { + case "d": + return "day" + case "w": + return "week" + case "m": + return "month" + case "y": + return "year" + default: + return "" + } +} + +func mapGLMRecencyFilter(rangeCode string) string { + switch rangeCode { + case "d": + return "oneDay" + case "w": + return "oneWeek" + case "m": + return "oneMonth" + case "y": + return "oneYear" + default: + return "noLimit" + } +} + +func mapBaiduRecencyFilter(rangeCode string) string { + switch rangeCode { + case "d", "w": + // Baidu does not expose a day-level filter. Use the closest supported + // window to keep recency bias instead of silently dropping the filter. + return "week" + case "m": + return "month" + case "y": + return "year" + default: + return "" + } } type BraveSearchProvider struct { @@ -95,9 +212,17 @@ type BraveSearchProvider struct { client *http.Client } -func (p *BraveSearchProvider) Search(ctx context.Context, query string, count int) (string, error) { +func (p *BraveSearchProvider) Search( + ctx context.Context, + query string, + count int, + rangeCode string, +) (string, error) { searchURL := fmt.Sprintf("https://api.search.brave.com/res/v1/web/search?q=%s&count=%d", url.QueryEscape(query), count) + if freshness := mapBraveFreshness(rangeCode); freshness != "" { + searchURL += "&freshness=" + url.QueryEscape(freshness) + } var lastErr error iter := p.keyPool.NewIterator() @@ -186,7 +311,12 @@ type TavilySearchProvider struct { client *http.Client } -func (p *TavilySearchProvider) Search(ctx context.Context, query string, count int) (string, error) { +func (p *TavilySearchProvider) Search( + ctx context.Context, + query string, + count int, + rangeCode string, +) (string, error) { searchURL := p.baseURL if searchURL == "" { searchURL = "https://api.tavily.com/search" @@ -210,6 +340,9 @@ func (p *TavilySearchProvider) Search(ctx context.Context, query string, count i "include_raw_content": false, "max_results": count, } + if timeRange := mapTavilyTimeRange(rangeCode); timeRange != "" { + payload["time_range"] = timeRange + } bodyBytes, err := json.Marshal(payload) if err != nil { @@ -289,8 +422,16 @@ type DuckDuckGoSearchProvider struct { client *http.Client } -func (p *DuckDuckGoSearchProvider) Search(ctx context.Context, query string, count int) (string, error) { +func (p *DuckDuckGoSearchProvider) Search( + ctx context.Context, + query string, + count int, + rangeCode string, +) (string, error) { searchURL := fmt.Sprintf("https://html.duckduckgo.com/html/?q=%s", url.QueryEscape(query)) + if dateFilter := mapDuckDuckGoDateFilter(rangeCode); dateFilter != "" { + searchURL += "&df=" + url.QueryEscape(dateFilter) + } req, err := http.NewRequestWithContext(ctx, "GET", searchURL, nil) if err != nil { @@ -313,7 +454,11 @@ func (p *DuckDuckGoSearchProvider) Search(ctx context.Context, query string, cou return p.extractResults(string(body), count, query) } -func (p *DuckDuckGoSearchProvider) extractResults(html string, count int, query string) (string, error) { +func (p *DuckDuckGoSearchProvider) extractResults( + html string, + count int, + query string, +) (string, error) { // Simple regex based extraction for DDG HTML // Strategy: Find all result containers or key anchors directly @@ -381,7 +526,12 @@ type PerplexitySearchProvider struct { client *http.Client } -func (p *PerplexitySearchProvider) Search(ctx context.Context, query string, count int) (string, error) { +func (p *PerplexitySearchProvider) Search( + ctx context.Context, + query string, + count int, + rangeCode string, +) (string, error) { searchURL := "https://api.perplexity.ai/chat/completions" var lastErr error @@ -401,19 +551,31 @@ func (p *PerplexitySearchProvider) Search(ctx context.Context, query string, cou "content": "You are a search assistant. Provide concise search results with titles, URLs, and brief descriptions in the following format:\n1. Title\n URL\n Description\n\nDo not add extra commentary.", }, { - "role": "user", - "content": fmt.Sprintf("Search for: %s. Provide up to %d relevant results.", query, count), + "role": "user", + "content": fmt.Sprintf( + "Search for: %s. Provide up to %d relevant results.", + query, + count, + ), }, }, "max_tokens": 1000, } + if recencyFilter := mapPerplexityRecencyFilter(rangeCode); recencyFilter != "" { + payload["search_recency_filter"] = recencyFilter + } payloadBytes, err := json.Marshal(payload) if err != nil { return "", fmt.Errorf("failed to marshal request: %w", err) } - req, err := http.NewRequestWithContext(ctx, "POST", searchURL, strings.NewReader(string(payloadBytes))) + req, err := http.NewRequestWithContext( + ctx, + "POST", + searchURL, + strings.NewReader(string(payloadBytes)), + ) if err != nil { return "", fmt.Errorf("failed to create request: %w", err) } @@ -463,7 +625,11 @@ func (p *PerplexitySearchProvider) Search(ctx context.Context, query string, cou return fmt.Sprintf("No results for: %s", query), nil } - return fmt.Sprintf("Results for: %s (via Perplexity)\n%s", query, searchResp.Choices[0].Message.Content), nil + return fmt.Sprintf( + "Results for: %s (via Perplexity)\n%s", + query, + searchResp.Choices[0].Message.Content, + ), nil } return "", fmt.Errorf("all api keys failed, last error: %w", lastErr) @@ -473,10 +639,18 @@ type SearXNGSearchProvider struct { baseURL string } -func (p *SearXNGSearchProvider) Search(ctx context.Context, query string, count int) (string, error) { +func (p *SearXNGSearchProvider) Search( + ctx context.Context, + query string, + count int, + rangeCode string, +) (string, error) { searchURL := fmt.Sprintf("%s/search?q=%s&format=json&categories=general", strings.TrimSuffix(p.baseURL, "/"), url.QueryEscape(query)) + if timeRange := mapSearXNGTimeRange(rangeCode); timeRange != "" { + searchURL += "&time_range=" + url.QueryEscape(timeRange) + } req, err := http.NewRequestWithContext(ctx, "GET", searchURL, nil) if err != nil { @@ -539,7 +713,12 @@ type GLMSearchProvider struct { client *http.Client } -func (p *GLMSearchProvider) Search(ctx context.Context, query string, count int) (string, error) { +func (p *GLMSearchProvider) Search( + ctx context.Context, + query string, + count int, + rangeCode string, +) (string, error) { searchURL := p.baseURL if searchURL == "" { searchURL = "https://open.bigmodel.cn/api/paas/v4/web_search" @@ -552,6 +731,9 @@ func (p *GLMSearchProvider) Search(ctx context.Context, query string, count int) "count": count, "content_size": "medium", } + if recencyFilter := mapGLMRecencyFilter(rangeCode); recencyFilter != "" { + payload["search_recency_filter"] = recencyFilter + } bodyBytes, err := json.Marshal(payload) if err != nil { @@ -620,7 +802,12 @@ type BaiduSearchProvider struct { client *http.Client } -func (p *BaiduSearchProvider) Search(ctx context.Context, query string, count int) (string, error) { +func (p *BaiduSearchProvider) Search( + ctx context.Context, + query string, + count int, + rangeCode string, +) (string, error) { searchURL := p.baseURL if searchURL == "" { searchURL = "https://qianfan.baidubce.com/v2/ai_search/web_search" @@ -636,6 +823,9 @@ func (p *BaiduSearchProvider) Search(ctx context.Context, query string, count in "search_source": "baidu_search_v2", "resource_type_filter": []map[string]any{{"type": "web", "top_k": count}}, } + if recencyFilter := mapBaiduRecencyFilter(rangeCode); recencyFilter != "" { + payload["search_recency_filter"] = recencyFilter + } bodyBytes, err := json.Marshal(payload) if err != nil { @@ -729,7 +919,7 @@ type WebSearchToolOptions struct { func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) { var provider SearchProvider - maxResults := 5 + maxResults := 10 // Priority: Perplexity > Brave > SearXNG > Tavily > DuckDuckGo > Baidu Search > GLM Search if opts.PerplexityEnabled && len(opts.PerplexityAPIKeys) > 0 { client, err := utils.CreateHTTPClient(opts.Proxy, perplexityTimeout) @@ -742,7 +932,7 @@ func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) { client: client, } if opts.PerplexityMaxResults > 0 { - maxResults = opts.PerplexityMaxResults + maxResults = min(opts.PerplexityMaxResults, 10) } } else if opts.BraveEnabled && len(opts.BraveAPIKeys) > 0 { client, err := utils.CreateHTTPClient(opts.Proxy, searchTimeout) @@ -751,12 +941,12 @@ func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) { } provider = &BraveSearchProvider{keyPool: NewAPIKeyPool(opts.BraveAPIKeys), proxy: opts.Proxy, client: client} if opts.BraveMaxResults > 0 { - maxResults = opts.BraveMaxResults + maxResults = min(opts.BraveMaxResults, 10) } } else if opts.SearXNGEnabled && opts.SearXNGBaseURL != "" { provider = &SearXNGSearchProvider{baseURL: opts.SearXNGBaseURL} if opts.SearXNGMaxResults > 0 { - maxResults = opts.SearXNGMaxResults + maxResults = min(opts.SearXNGMaxResults, 10) } } else if opts.TavilyEnabled && len(opts.TavilyAPIKeys) > 0 { client, err := utils.CreateHTTPClient(opts.Proxy, searchTimeout) @@ -770,7 +960,7 @@ func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) { client: client, } if opts.TavilyMaxResults > 0 { - maxResults = opts.TavilyMaxResults + maxResults = min(opts.TavilyMaxResults, 10) } } else if opts.DuckDuckGoEnabled { client, err := utils.CreateHTTPClient(opts.Proxy, searchTimeout) @@ -779,7 +969,7 @@ func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) { } provider = &DuckDuckGoSearchProvider{proxy: opts.Proxy, client: client} if opts.DuckDuckGoMaxResults > 0 { - maxResults = opts.DuckDuckGoMaxResults + maxResults = min(opts.DuckDuckGoMaxResults, 10) } } else if opts.BaiduSearchEnabled && opts.BaiduSearchAPIKey != "" { client, err := utils.CreateHTTPClient(opts.Proxy, perplexityTimeout) @@ -793,7 +983,7 @@ func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) { client: client, } if opts.BaiduSearchMaxResults > 0 { - maxResults = opts.BaiduSearchMaxResults + maxResults = min(opts.BaiduSearchMaxResults, 10) } } else if opts.GLMSearchEnabled && opts.GLMSearchAPIKey != "" { client, err := utils.CreateHTTPClient(opts.Proxy, searchTimeout) @@ -812,7 +1002,7 @@ func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) { client: client, } if opts.GLMSearchMaxResults > 0 { - maxResults = opts.GLMSearchMaxResults + maxResults = min(opts.GLMSearchMaxResults, 10) } } else { return nil, nil @@ -829,7 +1019,7 @@ func (t *WebSearchTool) Name() string { } func (t *WebSearchTool) Description() string { - return "Search the web for current information. Returns titles, URLs, and snippets from search results." + return "Search the web for current information. Supports query, count, and an optional temporal range filter. Returns titles, URLs, and snippets from search results." } func (t *WebSearchTool) Parameters() map[string]any { @@ -842,10 +1032,15 @@ func (t *WebSearchTool) Parameters() map[string]any { }, "count": map[string]any{ "type": "integer", - "description": "Number of results (1-10)", + "description": "Number of results (default: 10, max: 10)", "minimum": 1.0, "maximum": 10.0, }, + "range": map[string]any{ + "type": "string", + "description": "Optional time filter: d (day), w (week), m (month), y (year)", + "enum": []string{"d", "w", "m", "y"}, + }, }, "required": []string{"query"}, } @@ -853,18 +1048,36 @@ func (t *WebSearchTool) Parameters() map[string]any { func (t *WebSearchTool) Execute(ctx context.Context, args map[string]any) *ToolResult { query, ok := args["query"].(string) - if !ok { + if !ok || strings.TrimSpace(query) == "" { return ErrorResult("query is required") } + query = strings.TrimSpace(query) + count64, err := getInt64Arg(args, "count", int64(t.maxResults)) + if err != nil { + return ErrorResult(err.Error()) + } count := t.maxResults - if c, ok := args["count"].(float64); ok { - if int(c) > 0 && int(c) <= 10 { - count = int(c) + if count64 > 0 && count64 <= 10 { + count = int(count64) + } + + rangeCode, err := normalizeSearchRange("") + if err != nil { + return ErrorResult(err.Error()) + } + if rawRange, exists := args["range"]; exists { + rangeStr, ok := rawRange.(string) + if !ok { + return ErrorResult("range must be a string") + } + rangeCode, err = normalizeSearchRange(rangeStr) + if err != nil { + return ErrorResult(err.Error()) } } - result, err := t.provider.Search(ctx, query, count) + result, err := t.provider.Search(ctx, query, count, rangeCode) if err != nil { return ErrorResult(fmt.Sprintf("search failed: %v", err)) } @@ -1038,7 +1251,12 @@ func (t *WebFetchTool) Execute(ctx context.Context, args map[string]any) *ToolRe if err != nil { var maxBytesErr *http.MaxBytesError if errors.As(err, &maxBytesErr) { - return ErrorResult(fmt.Sprintf("failed to read response: size exceeded %d bytes limit", t.fetchLimitBytes)) + return ErrorResult( + fmt.Sprintf( + "failed to read response: size exceeded %d bytes limit", + t.fetchLimitBytes, + ), + ) } return ErrorResult(err.Error()) } @@ -1088,7 +1306,11 @@ func (t *WebFetchTool) Execute(ctx context.Context, args map[string]any) *ToolRe // If the charset is not utf-8, we might have to convert the bodyStr // before passing it to the HTML/Markdown parser if strings.ToLower(charset) != "utf-8" { - logger.WarnCF("tool", "Note: the content is not in UTF-8", map[string]any{"charset": charset}) + logger.WarnCF( + "tool", + "Note: the content is not in UTF-8", + map[string]any{"charset": charset}, + ) } } @@ -1232,7 +1454,11 @@ func newSafeDialContext( continue } attempted++ - conn, err := dialer.DialContext(ctx, network, net.JoinHostPort(ipAddr.IP.String(), port)) + conn, err := dialer.DialContext( + ctx, + network, + net.JoinHostPort(ipAddr.IP.String(), port), + ) if err == nil { return conn, nil } @@ -1240,10 +1466,17 @@ func newSafeDialContext( } if attempted == 0 { - return nil, fmt.Errorf("all resolved addresses for %s are private, restricted, or not whitelisted", host) + return nil, fmt.Errorf( + "all resolved addresses for %s are private, restricted, or not whitelisted", + host, + ) } if lastErr != nil { - return nil, fmt.Errorf("failed connecting to public addresses for %s: %w", host, lastErr) + return nil, fmt.Errorf( + "failed connecting to public addresses for %s: %w", + host, + lastErr, + ) } return nil, fmt.Errorf("failed connecting to public addresses for %s", host) } diff --git a/pkg/tools/web_test.go b/pkg/tools/web_test.go index 98c763193..de6187cfa 100644 --- a/pkg/tools/web_test.go +++ b/pkg/tools/web_test.go @@ -426,6 +426,96 @@ func TestWebTool_WebSearch_MissingQuery(t *testing.T) { } } +func TestNormalizeSearchRange(t *testing.T) { + tests := []struct { + name string + input string + want string + wantErr bool + }{ + {name: "empty", input: "", want: ""}, + {name: "day", input: "d", want: "d"}, + {name: "week uppercase trimmed", input: " W ", want: "w"}, + {name: "month", input: "m", want: "m"}, + {name: "year", input: "y", want: "y"}, + {name: "invalid", input: "q", wantErr: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := normalizeSearchRange(tt.input) + if tt.wantErr { + if err == nil { + t.Fatalf("expected error, got nil") + } + return + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got != tt.want { + t.Fatalf("normalizeSearchRange(%q) = %q, want %q", tt.input, got, tt.want) + } + }) + } +} + +func TestSearchRangeMappings(t *testing.T) { + if got := mapBraveFreshness("d"); got != "pd" { + t.Fatalf("mapBraveFreshness(d) = %q, want pd", got) + } + if got := mapBraveFreshness("y"); got != "py" { + t.Fatalf("mapBraveFreshness(y) = %q, want py", got) + } + if got := mapTavilyTimeRange("w"); got != "week" { + t.Fatalf("mapTavilyTimeRange(w) = %q, want week", got) + } + if got := mapPerplexityRecencyFilter("m"); got != "month" { + t.Fatalf("mapPerplexityRecencyFilter(m) = %q, want month", got) + } + if got := mapDuckDuckGoDateFilter("y"); got != "t" { + t.Fatalf("mapDuckDuckGoDateFilter(y) = %q, want t", got) + } + if got := mapSearXNGTimeRange("d"); got != "day" { + t.Fatalf("mapSearXNGTimeRange(d) = %q, want day", got) + } + if got := mapGLMRecencyFilter("w"); got != "oneWeek" { + t.Fatalf("mapGLMRecencyFilter(w) = %q, want oneWeek", got) + } + if got := mapGLMRecencyFilter(""); got != "noLimit" { + t.Fatalf("mapGLMRecencyFilter(\"\") = %q, want noLimit", got) + } + if got := mapBaiduRecencyFilter("d"); got != "week" { + t.Fatalf("mapBaiduRecencyFilter(d) = %q, want week", got) + } + if got := mapBaiduRecencyFilter("m"); got != "month" { + t.Fatalf("mapBaiduRecencyFilter(m) = %q, want month", got) + } +} + +func TestWebTool_WebSearch_InvalidRange(t *testing.T) { + tool, err := NewWebSearchTool(WebSearchToolOptions{ + BraveEnabled: true, + BraveAPIKeys: []string{"test-key"}, + BraveMaxResults: 5, + }) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + result := tool.Execute(context.Background(), map[string]any{ + "query": "test query", + "range": "invalid", + }) + + if !result.IsError { + t.Fatalf("expected invalid range to return error") + } + if !strings.Contains(result.ForLLM, "range must be one of: d, w, m, y") { + t.Fatalf("unexpected error message: %q", result.ForLLM) + } +} + // TestWebTool_WebFetch_HTMLExtraction verifies HTML text extraction func TestWebTool_WebFetch_HTMLExtraction(t *testing.T) { withPrivateWebFetchHostsAllowed(t) @@ -1069,6 +1159,45 @@ func TestWebTool_TavilySearch_Success(t *testing.T) { } } +func TestWebTool_TavilySearch_RangeMapping(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var payload map[string]any + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { + t.Fatalf("failed to decode payload: %v", err) + } + if payload["time_range"] != "week" { + t.Fatalf("expected time_range=week, got %v", payload["time_range"]) + } + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + json.NewEncoder(w).Encode(map[string]any{ + "results": []map[string]any{ + {"title": "Recent result", "url": "https://example.com/recent", "content": "snippet"}, + }, + }) + })) + defer server.Close() + + tool, err := NewWebSearchTool(WebSearchToolOptions{ + TavilyEnabled: true, + TavilyAPIKeys: []string{"test-key"}, + TavilyBaseURL: server.URL, + TavilyMaxResults: 5, + }) + if err != nil { + t.Fatalf("NewWebSearchTool() error: %v", err) + } + + result := tool.Execute(context.Background(), map[string]any{ + "query": "test query", + "range": "w", + }) + if result.IsError { + t.Fatalf("expected success, got %s", result.ForLLM) + } +} + // TestWebFetchTool_CloudflareChallenge_RetryWithHonestUA verifies that a 403 response // with cf-mitigated: challenge triggers a retry using the honest picoclaw User-Agent, // and that the retry response is returned when it succeeds. @@ -1297,6 +1426,39 @@ func TestWebTool_TavilySearch_Failover(t *testing.T) { } } +func TestWebTool_SearXNGSearch_RangeMapping(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got := r.URL.Query().Get("time_range"); got != "year" { + t.Fatalf("expected time_range=year, got %q", got) + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + json.NewEncoder(w).Encode(map[string]any{ + "results": []map[string]any{ + {"title": "Recent result", "url": "https://example.com/1", "content": "snippet"}, + }, + }) + })) + defer server.Close() + + tool, err := NewWebSearchTool(WebSearchToolOptions{ + SearXNGEnabled: true, + SearXNGBaseURL: server.URL, + SearXNGMaxResults: 5, + }) + if err != nil { + t.Fatalf("NewWebSearchTool() error: %v", err) + } + + result := tool.Execute(context.Background(), map[string]any{ + "query": "test query", + "range": "y", + }) + if result.IsError { + t.Fatalf("expected success, got %s", result.ForLLM) + } +} + func TestWebTool_GLMSearch_Success(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.Method != "POST" { @@ -1365,6 +1527,83 @@ func TestWebTool_GLMSearch_Success(t *testing.T) { } } +func TestWebTool_GLMSearch_RangeMapping(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var payload map[string]any + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { + t.Fatalf("failed to decode payload: %v", err) + } + if payload["search_recency_filter"] != "oneMonth" { + t.Fatalf("expected search_recency_filter=oneMonth, got %v", payload["search_recency_filter"]) + } + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + json.NewEncoder(w).Encode(map[string]any{ + "search_result": []map[string]any{ + {"title": "Recent GLM Result", "content": "snippet", "link": "https://example.com/glm-range"}, + }, + }) + })) + defer server.Close() + + tool, err := NewWebSearchTool(WebSearchToolOptions{ + GLMSearchEnabled: true, + GLMSearchAPIKey: "test-glm-key", + GLMSearchBaseURL: server.URL, + GLMSearchEngine: "search_std", + }) + if err != nil { + t.Fatalf("NewWebSearchTool() error: %v", err) + } + + result := tool.Execute(context.Background(), map[string]any{ + "query": "test query", + "range": "m", + }) + if result.IsError { + t.Fatalf("expected success, got %s", result.ForLLM) + } +} + +func TestWebTool_BaiduSearch_RangeMapping(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var payload map[string]any + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { + t.Fatalf("failed to decode payload: %v", err) + } + if payload["search_recency_filter"] != "week" { + t.Fatalf("expected search_recency_filter=week for day fallback, got %v", payload["search_recency_filter"]) + } + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + json.NewEncoder(w).Encode(map[string]any{ + "references": []map[string]any{ + {"title": "Recent Baidu Result", "url": "https://example.com/baidu", "content": "snippet"}, + }, + }) + })) + defer server.Close() + + tool, err := NewWebSearchTool(WebSearchToolOptions{ + BaiduSearchEnabled: true, + BaiduSearchAPIKey: "test-baidu-key", + BaiduSearchBaseURL: server.URL, + }) + if err != nil { + t.Fatalf("NewWebSearchTool() error: %v", err) + } + + result := tool.Execute(context.Background(), map[string]any{ + "query": "test query", + "range": "d", + }) + if result.IsError { + t.Fatalf("expected success, got %s", result.ForLLM) + } +} + func TestWebTool_GLMSearch_APIError(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusUnauthorized) diff --git a/web/backend/api/channels.go b/web/backend/api/channels.go index 21624d3ef..dd4c9af3d 100644 --- a/web/backend/api/channels.go +++ b/web/backend/api/channels.go @@ -22,8 +22,6 @@ var channelCatalog = []channelCatalogItem{ {Name: "qq", ConfigKey: "qq"}, {Name: "onebot", ConfigKey: "onebot"}, {Name: "wecom", ConfigKey: "wecom"}, - {Name: "wecom_app", ConfigKey: "wecom_app"}, - {Name: "wecom_aibot", ConfigKey: "wecom_aibot"}, {Name: "whatsapp", ConfigKey: "whatsapp", Variant: "bridge"}, {Name: "whatsapp_native", ConfigKey: "whatsapp", Variant: "native"}, {Name: "pico", ConfigKey: "pico"}, diff --git a/web/backend/api/config.go b/web/backend/api/config.go index e67e3e6d7..618b8438d 100644 --- a/web/backend/api/config.go +++ b/web/backend/api/config.go @@ -6,6 +6,7 @@ import ( "io" "net/http" "regexp" + "strings" "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/logger" @@ -16,6 +17,7 @@ func (h *Handler) registerConfigRoutes(mux *http.ServeMux) { mux.HandleFunc("GET /api/config", h.handleGetConfig) mux.HandleFunc("PUT /api/config", h.handleUpdateConfig) mux.HandleFunc("PATCH /api/config", h.handlePatchConfig) + mux.HandleFunc("POST /api/config/test-command-patterns", h.handleTestCommandPatterns) } // handleGetConfig returns the complete system configuration. @@ -179,6 +181,70 @@ func (h *Handler) handlePatchConfig(w http.ResponseWriter, r *http.Request) { json.NewEncoder(w).Encode(map[string]string{"status": "ok"}) } +// handleTestCommandPatterns tests a command against whitelist and blacklist patterns. +// +// POST /api/config/test-command-patterns +func (h *Handler) handleTestCommandPatterns(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(io.LimitReader(r.Body, 1<<20)) + if err != nil { + http.Error(w, "Failed to read request body", http.StatusBadRequest) + return + } + defer r.Body.Close() + + var req struct { + AllowPatterns []string `json:"allow_patterns"` + DenyPatterns []string `json:"deny_patterns"` + Command string `json:"command"` + } + if err := json.Unmarshal(body, &req); err != nil { + http.Error(w, fmt.Sprintf("Invalid JSON: %v", err), http.StatusBadRequest) + return + } + + lower := strings.ToLower(strings.TrimSpace(req.Command)) + + type result struct { + Allowed bool `json:"allowed"` + Blocked bool `json:"blocked"` + MatchedWhitelist *string `json:"matched_whitelist,omitempty"` + MatchedBlacklist *string `json:"matched_blacklist,omitempty"` + } + + resp := result{Allowed: false, Blocked: false} + + // Check whitelist first + for _, pattern := range req.AllowPatterns { + re, err := regexp.Compile(pattern) + if err != nil { + continue // skip invalid patterns + } + if re.MatchString(lower) { + resp.Allowed = true + resp.MatchedWhitelist = &pattern + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(resp) + return + } + } + + // Check blacklist + for _, pattern := range req.DenyPatterns { + re, err := regexp.Compile(pattern) + if err != nil { + continue + } + if re.MatchString(lower) { + resp.Blocked = true + resp.MatchedBlacklist = &pattern + break + } + } + + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(resp) +} + // validateConfig checks the config for common errors before saving. // Returns a list of human-readable error strings; empty means valid. func validateConfig(cfg *config.Config) []string { @@ -209,6 +275,15 @@ func validateConfig(cfg *config.Config) []string { errs = append(errs, "channels.discord.token is required when discord channel is enabled") } + if cfg.Channels.WeCom.Enabled { + if cfg.Channels.WeCom.BotID == "" { + errs = append(errs, "channels.wecom.bot_id is required when wecom channel is enabled") + } + if cfg.Channels.WeCom.Secret() == "" { + errs = append(errs, "channels.wecom.secret is required when wecom channel is enabled") + } + } + if cfg.Tools.Exec.Enabled { if cfg.Tools.Exec.EnableDenyPatterns { errs = append( diff --git a/web/backend/api/config_test.go b/web/backend/api/config_test.go index 9b05546f9..36acd95b0 100644 --- a/web/backend/api/config_test.go +++ b/web/backend/api/config_test.go @@ -282,3 +282,170 @@ func TestHandlePatchConfig_AllowsInvalidDenyRegexPatternsWhenDenyPatternsDisable t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String()) } } + +// testCommandPatterns is a helper that sets up a handler and sends a test-command-patterns request. +func testCommandPatterns(t *testing.T, configPath string, body string) *httptest.ResponseRecorder { + t.Helper() + h := NewHandler(configPath) + mux := http.NewServeMux() + h.RegisterRoutes(mux) + req := httptest.NewRequest(http.MethodPost, "/api/config/test-command-patterns", bytes.NewBufferString(body)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + mux.ServeHTTP(rec, req) + return rec +} + +func TestHandleTestCommandPatterns_MatchesWhitelist(t *testing.T) { + configPath, cleanup := setupOAuthTestEnv(t) + defer cleanup() + + rec := testCommandPatterns(t, configPath, `{ + "allow_patterns": ["^echo\\s+hello"], + "deny_patterns": ["^rm\\s+-rf"], + "command": "echo hello world" + }`) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String()) + } + if !bytes.Contains(rec.Body.Bytes(), []byte(`"allowed":true`)) { + t.Fatalf("expected allowed=true, body=%s", rec.Body.String()) + } + if bytes.Contains(rec.Body.Bytes(), []byte(`"blocked":true`)) { + t.Fatalf("expected blocked=false when whitelist matches, body=%s", rec.Body.String()) + } +} + +func TestHandleTestCommandPatterns_MatchesBlacklistNotWhitelist(t *testing.T) { + configPath, cleanup := setupOAuthTestEnv(t) + defer cleanup() + + rec := testCommandPatterns(t, configPath, `{ + "allow_patterns": ["^echo\\s+hello"], + "deny_patterns": ["^rm\\s+-rf"], + "command": "rm -rf /tmp" + }`) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String()) + } + if !bytes.Contains(rec.Body.Bytes(), []byte(`"blocked":true`)) { + t.Fatalf("expected blocked=true, body=%s", rec.Body.String()) + } + if bytes.Contains(rec.Body.Bytes(), []byte(`"allowed":true`)) { + t.Fatalf("expected allowed=false when blacklist matches but not whitelist, body=%s", rec.Body.String()) + } +} + +func TestHandleTestCommandPatterns_MatchesNeither(t *testing.T) { + configPath, cleanup := setupOAuthTestEnv(t) + defer cleanup() + + rec := testCommandPatterns(t, configPath, `{ + "allow_patterns": ["^echo\\s+hello"], + "deny_patterns": ["^rm\\s+-rf"], + "command": "ls -la" + }`) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String()) + } + if bytes.Contains(rec.Body.Bytes(), []byte(`"allowed":true`)) { + t.Fatalf("expected allowed=false, body=%s", rec.Body.String()) + } + if bytes.Contains(rec.Body.Bytes(), []byte(`"blocked":true`)) { + t.Fatalf("expected blocked=false, body=%s", rec.Body.String()) + } +} + +func TestHandleTestCommandPatterns_CaseInsensitiveWithGoFlag(t *testing.T) { + configPath, cleanup := setupOAuthTestEnv(t) + defer cleanup() + + rec := testCommandPatterns(t, configPath, `{ + "allow_patterns": ["(?i)^ECHO"], + "deny_patterns": [], + "command": "echo hello" + }`) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String()) + } + if !bytes.Contains(rec.Body.Bytes(), []byte(`"allowed":true`)) { + t.Fatalf("expected allowed=true with Go (?i) flag, body=%s", rec.Body.String()) + } +} + +func TestHandleTestCommandPatterns_EmptyPatterns(t *testing.T) { + configPath, cleanup := setupOAuthTestEnv(t) + defer cleanup() + + rec := testCommandPatterns(t, configPath, `{ + "allow_patterns": [], + "deny_patterns": [], + "command": "rm -rf /tmp" + }`) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String()) + } + if bytes.Contains(rec.Body.Bytes(), []byte(`"allowed":true`)) { + t.Fatalf("expected allowed=false with empty patterns, body=%s", rec.Body.String()) + } + if bytes.Contains(rec.Body.Bytes(), []byte(`"blocked":true`)) { + t.Fatalf("expected blocked=false with empty patterns, body=%s", rec.Body.String()) + } +} + +func TestHandleTestCommandPatterns_InvalidRegexSkipped(t *testing.T) { + configPath, cleanup := setupOAuthTestEnv(t) + defer cleanup() + + rec := testCommandPatterns(t, configPath, `{ + "allow_patterns": ["([[", "^echo"], + "deny_patterns": [], + "command": "echo hello" + }`) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String()) + } + if !bytes.Contains(rec.Body.Bytes(), []byte(`"allowed":true`)) { + t.Fatalf("expected allowed=true, invalid pattern skipped and valid one matched, body=%s", rec.Body.String()) + } +} + +func TestHandleTestCommandPatterns_ReturnsMatchedPattern(t *testing.T) { + configPath, cleanup := setupOAuthTestEnv(t) + defer cleanup() + + rec := testCommandPatterns(t, configPath, `{ + "allow_patterns": [], + "deny_patterns": ["\\$(?i)[a-zA-Z_]*(SECRET|KEY|PASSWORD|TOKEN|AUTH)[a-zA-Z0-9_]*"], + "command": "echo $GITHUB_API_KEY" + }`) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String()) + } + if !bytes.Contains(rec.Body.Bytes(), []byte(`"blocked":true`)) { + t.Fatalf("expected blocked=true, body=%s", rec.Body.String()) + } + if !bytes.Contains(rec.Body.Bytes(), []byte(`matched_blacklist`)) { + t.Fatalf("expected matched_blacklist field, body=%s", rec.Body.String()) + } +} + +func TestHandleTestCommandPatterns_InvalidJSON(t *testing.T) { + configPath, cleanup := setupOAuthTestEnv(t) + defer cleanup() + + h := NewHandler(configPath) + mux := http.NewServeMux() + h.RegisterRoutes(mux) + req := httptest.NewRequest( + http.MethodPost, + "/api/config/test-command-patterns", + bytes.NewBufferString(`{invalid json}`), + ) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + mux.ServeHTTP(rec, req) + if rec.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusBadRequest, rec.Body.String()) + } +} diff --git a/web/backend/api/models.go b/web/backend/api/models.go index 48babd8cd..38a55948b 100644 --- a/web/backend/api/models.go +++ b/web/backend/api/models.go @@ -42,6 +42,7 @@ type modelResponse struct { // Meta Configured bool `json:"configured"` IsDefault bool `json:"is_default"` + IsVirtual bool `json:"is_virtual"` } // handleListModels returns all model_list entries with masked API keys. @@ -86,6 +87,7 @@ func (h *Handler) handleListModels(w http.ResponseWriter, r *http.Request) { ExtraBody: m.ExtraBody, Configured: configured[i], IsDefault: m.ModelName == defaultModel, + IsVirtual: m.IsVirtual(), }) } @@ -202,8 +204,13 @@ func (h *Handler) handleUpdateModel(w http.ResponseWriter, r *http.Request) { } else { mc.ModelConfig.SetAPIKey(mc.APIKey) } + // Preserve existing ExtraBody when omitted (nil), but clear it when + // the frontend sends an empty object {} to indicate the field should + // be removed. if mc.ExtraBody == nil { mc.ExtraBody = cfg.ModelList[idx].ExtraBody + } else if len(mc.ExtraBody) == 0 { + mc.ExtraBody = nil } cfg.ModelList[idx] = &mc.ModelConfig @@ -288,11 +295,13 @@ func (h *Handler) handleSetDefaultModel(w http.ResponseWriter, r *http.Request) return } - // Verify the model_name exists in model_list + // Verify the model_name exists in model_list and is not a virtual model found := false + isVirtual := false for _, m := range cfg.ModelList { if m.ModelName == req.ModelName { found = true + isVirtual = m.IsVirtual() break } } @@ -300,6 +309,10 @@ func (h *Handler) handleSetDefaultModel(w http.ResponseWriter, r *http.Request) http.Error(w, fmt.Sprintf("Model %q not found in model_list", req.ModelName), http.StatusNotFound) return } + if isVirtual { + http.Error(w, fmt.Sprintf("Cannot set virtual model %q as default", req.ModelName), http.StatusBadRequest) + return + } cfg.Agents.Defaults.ModelName = req.ModelName diff --git a/web/backend/api/models_test.go b/web/backend/api/models_test.go index 9d3e72bd3..c80527fe3 100644 --- a/web/backend/api/models_test.go +++ b/web/backend/api/models_test.go @@ -356,6 +356,46 @@ func TestHandleAddModel_PersistsAPIKey(t *testing.T) { } } +// TestHandleSetDefaultModel_RejectsNonexistentModel tests that setting a non-existent +// model as default returns 404. This covers the case where virtual models (which are +// filtered by SaveConfig) cannot be set as default. +func TestHandleSetDefaultModel_RejectsNonexistentModel(t *testing.T) { + configPath, cleanup := setupOAuthTestEnv(t) + defer cleanup() + + // First save a valid config with a primary model + cfg, err := config.LoadConfig(configPath) + if err != nil { + t.Fatalf("LoadConfig() error = %v", err) + } + cfg.ModelList = []*config.ModelConfig{ + {ModelName: "gpt-4", Model: "openai/gpt-4o"}, + } + if err := config.SaveConfig(configPath, cfg); err != nil { + t.Fatalf("SaveConfig() error = %v", err) + } + + // Try to set a non-existent model (like a virtual model name) as default + h := NewHandler(configPath) + mux := http.NewServeMux() + h.RegisterRoutes(mux) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/models/default", bytes.NewBufferString(`{ + "model_name": "gpt-4__key_1" + }`)) + req.Header.Set("Content-Type", "application/json") + mux.ServeHTTP(rec, req) + + // Should return 404 because the virtual model doesn't exist in the persisted config + if rec.Code != http.StatusNotFound { + t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusNotFound, rec.Body.String()) + } + if !strings.Contains(rec.Body.String(), "not found") { + t.Fatalf("error message should mention 'not found', got: %s", rec.Body.String()) + } +} + func TestMaskAPIKey(t *testing.T) { tests := []struct { name string diff --git a/web/backend/api/router.go b/web/backend/api/router.go index d09f68eac..ce652d4c4 100644 --- a/web/backend/api/router.go +++ b/web/backend/api/router.go @@ -19,6 +19,8 @@ type Handler struct { oauthState map[string]string weixinMu sync.Mutex weixinFlows map[string]*weixinFlow + wecomMu sync.Mutex + wecomFlows map[string]*wecomFlow } // NewHandler creates an instance of the API handler. @@ -29,6 +31,7 @@ func NewHandler(configPath string) *Handler { oauthFlows: make(map[string]*oauthFlow), oauthState: make(map[string]string), weixinFlows: make(map[string]*weixinFlow), + wecomFlows: make(map[string]*wecomFlow), } } @@ -75,6 +78,9 @@ func (h *Handler) RegisterRoutes(mux *http.ServeMux) { // WeChat QR login flow h.registerWeixinRoutes(mux) + + // WeCom QR login flow + h.registerWecomRoutes(mux) } // Shutdown gracefully shuts down the handler, stopping the gateway if it was started by this handler. diff --git a/web/backend/api/wecom.go b/web/backend/api/wecom.go new file mode 100644 index 000000000..7dcec9f49 --- /dev/null +++ b/web/backend/api/wecom.go @@ -0,0 +1,424 @@ +package api + +import ( + "context" + "crypto/rand" + "encoding/hex" + "encoding/json" + "fmt" + "io" + "net/http" + "net/url" + "runtime" + "strconv" + "strings" + "time" + + "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/logger" +) + +const ( + wecomFlowTTL = 5 * time.Minute + wecomFlowGCAge = 30 * time.Minute + wecomQRSourceID = "picoclaw" + wecomQRGenerateEndpoint = "https://work.weixin.qq.com/ai/qc/generate" + wecomQRQueryEndpoint = "https://work.weixin.qq.com/ai/qc/query_result" + wecomQRHTTPTimeout = 15 * time.Second + wecomDefaultWebSocketURL = "wss://openws.work.weixin.qq.com" + wecomPollStartTimeout = 15 * time.Second + wecomPollStatusTimeout = 10 * time.Second +) + +const ( + wecomStatusWait = "wait" + wecomStatusScanned = "scaned" + wecomStatusConfirmed = "confirmed" + wecomStatusExpired = "expired" + wecomStatusError = "error" +) + +type wecomFlow struct { + ID string + SCode string + QRDataURI string + BotID string + Status string + Error string + CreatedAt time.Time + UpdatedAt time.Time + ExpiresAt time.Time +} + +type wecomFlowResponse struct { + FlowID string `json:"flow_id"` + Status string `json:"status"` + QRDataURI string `json:"qr_data_uri,omitempty"` + BotID string `json:"bot_id,omitempty"` + Error string `json:"error,omitempty"` +} + +type wecomQRGenerateResponse struct { + ErrCode int `json:"errcode,omitempty"` + ErrMsg string `json:"errmsg,omitempty"` + Data struct { + SCode string `json:"scode"` + AuthURL string `json:"auth_url"` + } `json:"data"` +} + +type wecomQRQueryResponse struct { + ErrCode int `json:"errcode,omitempty"` + ErrMsg string `json:"errmsg,omitempty"` + Data struct { + Status string `json:"status"` + BotInfo struct { + BotID string `json:"botid"` + Secret string `json:"secret"` + } `json:"bot_info"` + } `json:"data"` +} + +// registerWecomRoutes binds WeCom QR login endpoints to the ServeMux. +func (h *Handler) registerWecomRoutes(mux *http.ServeMux) { + mux.HandleFunc("POST /api/wecom/flows", h.handleStartWecomFlow) + mux.HandleFunc("GET /api/wecom/flows/{id}", h.handlePollWecomFlow) +} + +// handleStartWecomFlow starts a new WeCom QR login flow. +// +// POST /api/wecom/flows +func (h *Handler) handleStartWecomFlow(w http.ResponseWriter, r *http.Request) { + ctx, cancel := context.WithTimeout(r.Context(), wecomPollStartTimeout) + defer cancel() + + session, err := fetchWecomQRCode(ctx) + if err != nil { + http.Error(w, fmt.Sprintf("failed to get QR code: %v", err), http.StatusInternalServerError) + return + } + + dataURI, err := generateQRDataURI(session.Data.AuthURL) + if err != nil { + http.Error(w, fmt.Sprintf("failed to generate QR image: %v", err), http.StatusInternalServerError) + return + } + + now := time.Now() + flow := &wecomFlow{ + ID: newWecomFlowID(), + SCode: session.Data.SCode, + QRDataURI: dataURI, + Status: wecomStatusWait, + CreatedAt: now, + UpdatedAt: now, + ExpiresAt: now.Add(wecomFlowTTL), + } + h.storeWecomFlow(flow) + + logger.InfoCF("wecom", "QR flow started", map[string]any{"flow_id": flow.ID}) + + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(wecomFlowResponse{ + FlowID: flow.ID, + Status: flow.Status, + QRDataURI: flow.QRDataURI, + }) +} + +// handlePollWecomFlow polls the WeCom API for QR code status and updates the flow. +// +// GET /api/wecom/flows/{id} +func (h *Handler) handlePollWecomFlow(w http.ResponseWriter, r *http.Request) { + flowID := strings.TrimSpace(r.PathValue("id")) + if flowID == "" { + http.Error(w, "missing flow id", http.StatusBadRequest) + return + } + + flow, ok := h.getWecomFlow(flowID) + if !ok { + http.Error(w, "flow not found", http.StatusNotFound) + return + } + + if flow.Status == wecomStatusConfirmed || + flow.Status == wecomStatusExpired || + flow.Status == wecomStatusError { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(wecomFlowResponse{ + FlowID: flow.ID, + Status: flow.Status, + BotID: flow.BotID, + Error: flow.Error, + }) + return + } + + ctx, cancel := context.WithTimeout(r.Context(), wecomPollStatusTimeout) + defer cancel() + + statusResp, err := queryWecomQRCodeStatus(ctx, flow.SCode) + if err != nil { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(wecomFlowResponse{ + FlowID: flow.ID, + Status: flow.Status, + QRDataURI: flow.QRDataURI, + }) + return + } + + switch strings.ToLower(statusResp.Data.Status) { + case wecomStatusWait: + // no-op + case wecomStatusScanned, "scanned": + h.updateWecomFlowStatus(flowID, wecomStatusScanned) + case "success": + if statusResp.Data.BotInfo.BotID == "" || statusResp.Data.BotInfo.Secret == "" { + h.setWecomFlowError(flowID, "login confirmed but missing bot credentials") + break + } + if saveErr := h.saveWecomBinding( + statusResp.Data.BotInfo.BotID, + statusResp.Data.BotInfo.Secret, + ); saveErr != nil { + h.setWecomFlowError(flowID, fmt.Sprintf("failed to save credentials: %v", saveErr)) + logger.ErrorCF("wecom", "failed to save credentials", map[string]any{"error": saveErr.Error()}) + break + } + h.setWecomFlowConfirmed(flowID, statusResp.Data.BotInfo.BotID) + logger.InfoCF("wecom", "QR login confirmed, credentials saved", map[string]any{ + "flow_id": flowID, + "bot_id": statusResp.Data.BotInfo.BotID, + }) + case wecomStatusExpired: + h.updateWecomFlowStatus(flowID, wecomStatusExpired) + } + + flow, _ = h.getWecomFlow(flowID) + w.Header().Set("Content-Type", "application/json") + resp := wecomFlowResponse{ + FlowID: flow.ID, + Status: flow.Status, + BotID: flow.BotID, + Error: flow.Error, + } + if flow.Status == wecomStatusWait || flow.Status == wecomStatusScanned { + resp.QRDataURI = flow.QRDataURI + } + _ = json.NewEncoder(w).Encode(resp) +} + +func (h *Handler) saveWecomBinding(botID, secret string) error { + cfg, err := config.LoadConfig(h.configPath) + if err != nil { + return fmt.Errorf("load config: %w", err) + } + + cfg.Channels.WeCom.Enabled = true + cfg.Channels.WeCom.BotID = botID + cfg.Channels.WeCom.SetSecret(secret) + if strings.TrimSpace(cfg.Channels.WeCom.WebSocketURL) == "" { + cfg.Channels.WeCom.WebSocketURL = wecomDefaultWebSocketURL + } + if err := config.SaveConfig(h.configPath, cfg); err != nil { + return err + } + + status := h.gatewayStatusData() + gatewayStatus, _ := status["gateway_status"].(string) + if gatewayStatus != "running" { + return nil + } + + if _, err := h.RestartGateway(); err != nil { + logger.ErrorCF("wecom", "failed to restart gateway after saving binding", map[string]any{ + "error": err.Error(), + }) + } + return nil +} + +func fetchWecomQRCode(ctx context.Context) (wecomQRGenerateResponse, error) { + targetURL, err := buildWecomQRGenerateURL(wecomQRGenerateEndpoint, wecomQRSourceID, wecomPlatformCode()) + if err != nil { + return wecomQRGenerateResponse{}, err + } + + var resp wecomQRGenerateResponse + if err := doWecomJSONGet(ctx, targetURL, &resp); err != nil { + return wecomQRGenerateResponse{}, err + } + if resp.ErrCode != 0 { + return wecomQRGenerateResponse{}, fmt.Errorf( + "errcode=%d errmsg=%s", + resp.ErrCode, + resp.ErrMsg, + ) + } + if resp.Data.SCode == "" || resp.Data.AuthURL == "" { + return wecomQRGenerateResponse{}, fmt.Errorf("response missing scode or auth_url") + } + return resp, nil +} + +func queryWecomQRCodeStatus(ctx context.Context, scode string) (wecomQRQueryResponse, error) { + targetURL, err := buildWecomQRQueryURL(wecomQRQueryEndpoint, scode) + if err != nil { + return wecomQRQueryResponse{}, err + } + + var resp wecomQRQueryResponse + if err := doWecomJSONGet(ctx, targetURL, &resp); err != nil { + return wecomQRQueryResponse{}, err + } + if resp.ErrCode != 0 { + return wecomQRQueryResponse{}, fmt.Errorf( + "errcode=%d errmsg=%s", + resp.ErrCode, + resp.ErrMsg, + ) + } + return resp, nil +} + +func buildWecomQRGenerateURL(baseURL, sourceID string, platformCode int) (string, error) { + u, err := url.Parse(baseURL) + if err != nil { + return "", fmt.Errorf("invalid WeCom QR generate URL: %w", err) + } + + query := u.Query() + query.Set("source", sourceID) + query.Set("sourceID", sourceID) + query.Set("plat", strconv.Itoa(platformCode)) + u.RawQuery = query.Encode() + + return u.String(), nil +} + +func buildWecomQRQueryURL(baseURL, scode string) (string, error) { + u, err := url.Parse(baseURL) + if err != nil { + return "", fmt.Errorf("invalid WeCom QR query URL: %w", err) + } + + query := u.Query() + query.Set("scode", scode) + u.RawQuery = query.Encode() + + return u.String(), nil +} + +func doWecomJSONGet(ctx context.Context, targetURL string, out any) error { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, targetURL, nil) + if err != nil { + return err + } + + client := &http.Client{Timeout: wecomQRHTTPTimeout} + resp, err := client.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + body, readErr := io.ReadAll(io.LimitReader(resp.Body, 8192)) + if readErr != nil { + return fmt.Errorf("unexpected status %s", resp.Status) + } + return fmt.Errorf("unexpected status %s: %s", resp.Status, strings.TrimSpace(string(body))) + } + + if err := json.NewDecoder(resp.Body).Decode(out); err != nil { + return fmt.Errorf("decode JSON response: %w", err) + } + return nil +} + +func wecomPlatformCode() int { + switch runtime.GOOS { + case "darwin": + return 1 + case "windows": + return 2 + case "linux": + return 3 + default: + return 0 + } +} + +func newWecomFlowID() string { + buf := make([]byte, 12) + if _, err := rand.Read(buf); err != nil { + return fmt.Sprintf("wc_%d", time.Now().UnixNano()) + } + return "wc_" + hex.EncodeToString(buf) +} + +func (h *Handler) storeWecomFlow(flow *wecomFlow) { + h.wecomMu.Lock() + defer h.wecomMu.Unlock() + h.gcWecomFlowsLocked(time.Now()) + h.wecomFlows[flow.ID] = flow +} + +func (h *Handler) getWecomFlow(flowID string) (*wecomFlow, bool) { + h.wecomMu.Lock() + defer h.wecomMu.Unlock() + h.gcWecomFlowsLocked(time.Now()) + flow, ok := h.wecomFlows[flowID] + if !ok { + return nil, false + } + cp := *flow + return &cp, true +} + +func (h *Handler) updateWecomFlowStatus(flowID, status string) { + h.wecomMu.Lock() + defer h.wecomMu.Unlock() + if flow, ok := h.wecomFlows[flowID]; ok { + flow.Status = status + flow.UpdatedAt = time.Now() + } +} + +func (h *Handler) setWecomFlowConfirmed(flowID, botID string) { + h.wecomMu.Lock() + defer h.wecomMu.Unlock() + if flow, ok := h.wecomFlows[flowID]; ok { + flow.Status = wecomStatusConfirmed + flow.BotID = botID + flow.UpdatedAt = time.Now() + } +} + +func (h *Handler) setWecomFlowError(flowID, errMsg string) { + h.wecomMu.Lock() + defer h.wecomMu.Unlock() + if flow, ok := h.wecomFlows[flowID]; ok { + flow.Status = wecomStatusError + flow.Error = errMsg + flow.UpdatedAt = time.Now() + } +} + +func (h *Handler) gcWecomFlowsLocked(now time.Time) { + for id, flow := range h.wecomFlows { + if flow.Status == wecomStatusWait || flow.Status == wecomStatusScanned { + if !flow.ExpiresAt.IsZero() && now.After(flow.ExpiresAt) { + flow.Status = wecomStatusExpired + flow.UpdatedAt = now + } + } + if flow.Status != wecomStatusWait && + flow.Status != wecomStatusScanned && + now.Sub(flow.UpdatedAt) > wecomFlowGCAge { + delete(h.wecomFlows, id) + } + } +} diff --git a/web/backend/main.go b/web/backend/main.go index 2f181603e..6987a4515 100644 --- a/web/backend/main.go +++ b/web/backend/main.go @@ -98,7 +98,7 @@ func main() { defer logger.DisableFileLogging() } - logger.InfoC("web", "PicoClaw Launcher starting...") + logger.InfoC("web", fmt.Sprintf("%s Launcher %s starting...", appName, appVersion)) logger.InfoC("web", fmt.Sprintf("PicoClaw Home: %s", picoHome)) // Set language from command line or auto-detect diff --git a/web/backend/systray.go b/web/backend/systray.go index fde2e115e..9dcc025df 100644 --- a/web/backend/systray.go +++ b/web/backend/systray.go @@ -3,7 +3,6 @@ package main import ( - _ "embed" "fmt" "fyne.io/systray" @@ -93,8 +92,3 @@ func onReady() { func onExit() { logger.Info(T(Exiting)) } - -// getIcon returns the system tray icon -func getIcon() []byte { - return iconData -} diff --git a/web/backend/systray_icon_nonwindows.go b/web/backend/systray_icon_nonwindows.go new file mode 100644 index 000000000..0117a9ae8 --- /dev/null +++ b/web/backend/systray_icon_nonwindows.go @@ -0,0 +1,12 @@ +//go:build !windows && ((!darwin && !freebsd) || cgo) + +package main + +import _ "embed" + +//go:embed icon.png +var iconPNG []byte + +func getIcon() []byte { + return iconPNG +} diff --git a/web/backend/systray_windows.go b/web/backend/systray_icon_windows.go similarity index 53% rename from web/backend/systray_windows.go rename to web/backend/systray_icon_windows.go index cc1885155..c265e2f9c 100644 --- a/web/backend/systray_windows.go +++ b/web/backend/systray_icon_windows.go @@ -5,4 +5,8 @@ package main import _ "embed" //go:embed icon.ico -var iconData []byte +var iconICO []byte + +func getIcon() []byte { + return iconICO +} diff --git a/web/backend/tray_stub_nocgo.go b/web/backend/systray_stub_nocgo.go similarity index 88% rename from web/backend/tray_stub_nocgo.go rename to web/backend/systray_stub_nocgo.go index 13ecfd2cb..9e75e112a 100644 --- a/web/backend/tray_stub_nocgo.go +++ b/web/backend/systray_stub_nocgo.go @@ -13,6 +13,7 @@ import ( "github.com/sipeed/picoclaw/pkg/logger" ) +// runTray falls back to a headless mode on platforms where systray requires cgo. func runTray() { logger.Infof("System tray is unavailable in %s builds without cgo; running without tray", runtime.GOOS) diff --git a/web/backend/systray_unix.go b/web/backend/systray_unix.go deleted file mode 100644 index 0f9d2bb51..000000000 --- a/web/backend/systray_unix.go +++ /dev/null @@ -1,8 +0,0 @@ -//go:build !windows - -package main - -import _ "embed" - -//go:embed icon.png -var iconData []byte diff --git a/web/frontend/src/api/channels.ts b/web/frontend/src/api/channels.ts index d4c3ac74b..85550ca81 100644 --- a/web/frontend/src/api/channels.ts +++ b/web/frontend/src/api/channels.ts @@ -72,6 +72,14 @@ export interface WeixinFlowResponse { error?: string } +export interface WecomFlowResponse { + flow_id: string + status: "wait" | "scaned" | "confirmed" | "expired" | "error" + qr_data_uri?: string + bot_id?: string + error?: string +} + export async function startWeixinFlow(): Promise { return request("/api/weixin/flows", { method: "POST" }) } @@ -84,4 +92,16 @@ export async function pollWeixinFlow( ) } +export async function startWecomFlow(): Promise { + return request("/api/wecom/flows", { method: "POST" }) +} + +export async function pollWecomFlow( + flowID: string, +): Promise { + return request( + `/api/wecom/flows/${encodeURIComponent(flowID)}`, + ) +} + export type { ChannelsCatalogResponse, ConfigActionResponse } diff --git a/web/frontend/src/api/models.ts b/web/frontend/src/api/models.ts index 2fd042593..aa66a7389 100644 --- a/web/frontend/src/api/models.ts +++ b/web/frontend/src/api/models.ts @@ -21,6 +21,7 @@ export interface ModelInfo { // Meta configured: boolean is_default: boolean + is_virtual: boolean } interface ModelsListResponse { diff --git a/web/frontend/src/components/channels/channel-config-page.tsx b/web/frontend/src/components/channels/channel-config-page.tsx index ee483d652..6af821ac9 100644 --- a/web/frontend/src/components/channels/channel-config-page.tsx +++ b/web/frontend/src/components/channels/channel-config-page.tsx @@ -1,4 +1,4 @@ -import { IconLoader2 } from "@tabler/icons-react" +import { IconAlertTriangle, IconLoader2 } from "@tabler/icons-react" import { useCallback, useEffect, useMemo, useRef, useState } from "react" import { useTranslation } from "react-i18next" @@ -15,6 +15,7 @@ import { FeishuForm } from "@/components/channels/channel-forms/feishu-form" import { GenericForm } from "@/components/channels/channel-forms/generic-form" import { SlackForm } from "@/components/channels/channel-forms/slack-form" import { TelegramForm } from "@/components/channels/channel-forms/telegram-form" +import { WecomForm } from "@/components/channels/channel-forms/wecom-form" import { WeixinForm } from "@/components/channels/channel-forms/weixin-form" import { PageHeader } from "@/components/page-header" import { Button } from "@/components/ui/button" @@ -145,13 +146,7 @@ function isConfigured( case "weixin": return asString(config.account_id) !== "" case "wecom": - return asString(config.token) !== "" - case "wecom_app": - return ( - asString(config.corp_id) !== "" && asString(config.corp_secret) !== "" - ) - case "wecom_aibot": - return asString(config.token) !== "" + return asString(config.bot_id) !== "" case "whatsapp": return asString(config.bridge_url) !== "" case "whatsapp_native": @@ -192,11 +187,7 @@ function getRequiredFieldKeys(channelName: string): string[] { case "onebot": return ["ws_url"] case "wecom": - return ["token"] - case "wecom_app": - return ["corp_id", "corp_secret"] - case "wecom_aibot": - return ["token"] + return [] case "whatsapp": return ["bridge_url"] case "pico": @@ -336,6 +327,8 @@ export function ChannelConfigPage({ channelName }: ChannelConfigPageProps) { return getChannelDisplayName(channel, t) }, [channel, channelName, t]) + const hidesPageLevelEnableToggle = channel?.name === "wecom" + const hiddenKeys = useMemo(() => { if (!channel) return [] if (channel.name === "whatsapp") { @@ -420,6 +413,36 @@ export function ChannelConfigPage({ channelName }: ChannelConfigPageProps) { } }, [loadData, t]) + const handleWecomBindSuccess = useCallback(async () => { + try { + setEnabled(true) + await Promise.all([loadData(true), refreshGatewayState({ force: true })]) + } catch (e) { + const message = + e instanceof Error ? e.message : t("channels.page.saveError") + setServerError(message) + await loadData(true) + } + }, [loadData, t]) + + const handleWecomEnabledChange = useCallback( + async (nextEnabled: boolean) => { + try { + setEnabled(nextEnabled) + await Promise.all([ + loadData(true), + refreshGatewayState({ force: true }), + ]) + } catch (e) { + const message = + e instanceof Error ? e.message : t("channels.page.saveError") + setServerError(message) + await loadData(true) + } + }, + [loadData, t], + ) + const renderForm = () => { if (!channel) return null const isEdit = configured @@ -470,6 +493,27 @@ export function ChannelConfigPage({ channelName }: ChannelConfigPageProps) { onBindSuccess={() => void handleWeixinBindSuccess()} /> ) + case "wecom": + return ( + <> + void handleWecomBindSuccess()} + onEnabledChange={(nextEnabled) => + void handleWecomEnabledChange(nextEnabled) + } + /> + + + ) default: return ( -
-

- {t("channels.page.enableLabel")} -

- -
+ {channel?.name === "weixin" && ( +
+
+ +
+

+ {t("channels.weixin.warningTitle")} +

+

+ {t("channels.weixin.warningDesc")} +

+
+
+
+ )} + + {!hidesPageLevelEnableToggle && ( +
+

+ {t("channels.page.enableLabel")} +

+ +
+ )} {renderForm()} diff --git a/web/frontend/src/components/channels/channel-forms/generic-form.tsx b/web/frontend/src/components/channels/channel-forms/generic-form.tsx index db14fc206..936802944 100644 --- a/web/frontend/src/components/channels/channel-forms/generic-form.tsx +++ b/web/frontend/src/components/channels/channel-forms/generic-form.tsx @@ -28,6 +28,7 @@ const SECRET_FIELDS = new Set([ "encoding_aes_key", "encrypt_key", "verification_token", + "secret", "password", "nickserv_password", "sasl_password", @@ -44,6 +45,7 @@ const OBJECT_FIELDS = new Set([ "allow_token_query", "allow_from", "allow_origins", + "groups", ]) function formatLabel(key: string): string { @@ -118,6 +120,16 @@ export function GenericForm({ app_id: t("channels.form.desc.appId"), client_id: t("channels.form.desc.clientId"), corp_id: t("channels.form.desc.corpId"), + bot_id: t("channels.form.desc.appId"), + websocket_url: t("channels.form.desc.wsUrl"), + dm_policy: t("channels.form.desc.genericField", { field: "DM policy" }), + group_policy: t("channels.form.desc.genericField", { + field: "group policy", + }), + group_allow_from: t("channels.form.desc.allowFrom"), + send_thinking_message: t("channels.form.desc.genericField", { + field: "thinking message behavior", + }), agent_id: t("channels.form.desc.agentId"), webhook_url: t("channels.form.desc.webhookUrl"), webhook_host: t("channels.form.desc.webhookHost"), diff --git a/web/frontend/src/components/channels/channel-forms/wecom-form.tsx b/web/frontend/src/components/channels/channel-forms/wecom-form.tsx new file mode 100644 index 000000000..744c87ba2 --- /dev/null +++ b/web/frontend/src/components/channels/channel-forms/wecom-form.tsx @@ -0,0 +1,367 @@ +import { + IconCheck, + IconLoader2, + IconQrcode, + IconRefresh, + IconX, +} from "@tabler/icons-react" +import { useCallback, useEffect, useRef, useState } from "react" +import { useTranslation } from "react-i18next" + +import type { ChannelConfig } from "@/api/channels" +import { patchAppConfig, pollWecomFlow, startWecomFlow } from "@/api/channels" +import { Button } from "@/components/ui/button" +import { Switch } from "@/components/ui/switch" + +type BindingState = + | "idle" + | "loading" + | "waiting" + | "scaned" + | "confirmed" + | "expired" + | "error" + +interface WecomFormProps { + config: ChannelConfig + isEdit: boolean + onBindSuccess?: () => void + onEnabledChange?: (enabled: boolean) => void +} + +function asString(value: unknown): string { + return typeof value === "string" ? value : "" +} + +export function WecomForm({ + config, + isEdit, + onBindSuccess, + onEnabledChange, +}: WecomFormProps) { + const { t } = useTranslation() + + const [bindState, setBindState] = useState("idle") + const [qrDataURI, setQrDataURI] = useState(null) + const [botID, setBotID] = useState(null) + const [errorMsg, setErrorMsg] = useState("") + const [enabled, setEnabled] = useState(config.enabled === true) + const [toggleSaving, setToggleSaving] = useState(false) + const [toggleError, setToggleError] = useState("") + + const pollTimerRef = useRef | null>(null) + const pollGenerationRef = useRef(0) + const existingBotID = asString(config.bot_id) + const isBound = isEdit && existingBotID !== "" + + const stopPolling = useCallback(() => { + pollGenerationRef.current += 1 + if (pollTimerRef.current !== null) { + clearInterval(pollTimerRef.current) + pollTimerRef.current = null + } + }, []) + + useEffect(() => () => stopPolling(), [stopPolling]) + + useEffect(() => { + setEnabled(config.enabled === true) + }, [config.enabled]) + + useEffect(() => { + if (!existingBotID) return + stopPolling() + setBotID(existingBotID) + setBindState("confirmed") + setErrorMsg("") + }, [existingBotID, stopPolling]) + + const startPolling = useCallback( + (id: string) => { + stopPolling() + const generation = pollGenerationRef.current + let inFlight = false + pollTimerRef.current = setInterval(async () => { + if (inFlight) return + inFlight = true + try { + const resp = await pollWecomFlow(id) + if (generation !== pollGenerationRef.current) { + return + } + if (resp.status === "scaned") { + setBindState("scaned") + } else if (resp.status === "confirmed") { + stopPolling() + setBotID(resp.bot_id ?? existingBotID ?? null) + setBindState("confirmed") + onBindSuccess?.() + } else if (resp.status === "expired") { + stopPolling() + setBindState("expired") + } else if (resp.status === "error") { + stopPolling() + setBindState("error") + setErrorMsg(resp.error ?? t("channels.wecom.errorGeneric")) + } + } catch { + // transient network error — keep polling + } finally { + inFlight = false + } + }, 2000) + }, + [existingBotID, onBindSuccess, stopPolling, t], + ) + + const handleEnabledChange = useCallback( + async (checked: boolean) => { + if (!existingBotID || toggleSaving) { + return + } + setToggleSaving(true) + setToggleError("") + try { + await patchAppConfig({ + channels: { + wecom: { + enabled: checked, + }, + }, + }) + setEnabled(checked) + onEnabledChange?.(checked) + } catch (e) { + setToggleError( + e instanceof Error ? e.message : t("channels.wecom.errorGeneric"), + ) + } finally { + setToggleSaving(false) + } + }, + [existingBotID, onEnabledChange, t, toggleSaving], + ) + + const handleBind = async () => { + setBindState("loading") + setErrorMsg("") + setToggleError("") + setQrDataURI(null) + stopPolling() + try { + const resp = await startWecomFlow() + setQrDataURI(resp.qr_data_uri ?? null) + setBindState("waiting") + startPolling(resp.flow_id) + } catch (e) { + setBindState("error") + setErrorMsg( + e instanceof Error ? e.message : t("channels.wecom.errorGeneric"), + ) + } + } + + const handleRebind = () => { + stopPolling() + setBindState("idle") + setQrDataURI(null) + setBotID(null) + setErrorMsg("") + void handleBind() + } + + const renderBindSection = () => { + if (bindState === "idle") { + if (isBound) { + return ( +
+
+ + {t("channels.wecom.bound")} +
+ {existingBotID && ( +

+ {existingBotID} +

+ )} + +
+ ) + } + return ( +
+

+ {t("channels.wecom.notBound")} +

+ +
+ ) + } + + if (bindState === "loading") { + return ( +
+ +

+ {t("channels.wecom.generating")} +

+
+ ) + } + + if (bindState === "waiting" || bindState === "scaned") { + return ( +
+ {qrDataURI ? ( + WeCom QR Code + ) : ( +
+ +
+ )} + {bindState === "scaned" ? ( +
+ + {t("channels.wecom.scanned")} +
+ ) : ( +

+ {t("channels.wecom.scanHint")} +

+ )} + +
+ ) + } + + if (bindState === "confirmed") { + return ( +
+
+ +
+

+ {t("channels.wecom.bound")} +

+ {botID && ( +

{botID}

+ )} + +
+ ) + } + + if (bindState === "expired") { + return ( +
+
+ +
+

+ {t("channels.wecom.expired")} +

+ +
+ ) + } + + if (bindState === "error") { + return ( +
+
+ +
+

+ {errorMsg || t("channels.wecom.errorGeneric")} +

+ +
+ ) + } + + return null + } + + return ( +
+
+
+
+

+ {t("channels.page.enableLabel")} +

+

+ {isBound + ? t("channels.wecom.enableDesc") + : t("channels.wecom.enableBindFirst")} +

+
+ void handleEnabledChange(checked)} + /> +
+ {toggleError && ( +

{toggleError}

+ )} +
+ +
+
+

{t("channels.wecom.bindTitle")}

+

+ {t("channels.wecom.bindDesc")} +

+
+ {renderBindSection()} +
+
+ ) +} diff --git a/web/frontend/src/components/chat/user-message.tsx b/web/frontend/src/components/chat/user-message.tsx index 0035f14fa..2baee87f3 100644 --- a/web/frontend/src/components/chat/user-message.tsx +++ b/web/frontend/src/components/chat/user-message.tsx @@ -5,7 +5,7 @@ interface UserMessageProps { export function UserMessage({ content }: UserMessageProps) { return (
-
+
{content}
diff --git a/web/frontend/src/components/config/config-page.tsx b/web/frontend/src/components/config/config-page.tsx index 24a719d86..46f62f426 100644 --- a/web/frontend/src/components/config/config-page.tsx +++ b/web/frontend/src/components/config/config-page.tsx @@ -208,6 +208,7 @@ export function ConfigPage() { defaults: { workspace, restrict_to_workspace: form.restrictToWorkspace, + split_on_marker: form.splitOnMarker, tool_feedback: { enabled: form.toolFeedbackEnabled, max_args_length: toolFeedbackMaxArgsLength, diff --git a/web/frontend/src/components/config/config-sections.tsx b/web/frontend/src/components/config/config-sections.tsx index 5482b0a35..b5bec8acd 100644 --- a/web/frontend/src/components/config/config-sections.tsx +++ b/web/frontend/src/components/config/config-sections.tsx @@ -1,3 +1,4 @@ +import { useState } from "react" import type { ReactNode } from "react" import { useTranslation } from "react-i18next" @@ -7,6 +8,7 @@ import { type LauncherForm, } from "@/components/config/form-model" import { Field, SwitchCardField } from "@/components/shared-form" +import { Button } from "@/components/ui/button" import { Card, CardContent, @@ -93,6 +95,16 @@ export function AgentDefaultsSection({ } /> + + onFieldChange("splitOnMarker", checked) + } + /> + (null) + const [isLoading, setIsLoading] = useState(false) + + const testPatterns = async () => { + if (!testCommand.trim()) { + setTestResult(null) + return + } + + const allowPatterns = form.customAllowPatternsText + .split("\n") + .map((p) => p.trim()) + .filter((p) => p.length > 0) + const denyPatterns = form.enableDenyPatterns + ? form.customDenyPatternsText + .split("\n") + .map((p) => p.trim()) + .filter((p) => p.length > 0) + : [] + + setIsLoading(true) + try { + const res = await fetch("/api/config/test-command-patterns", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + allow_patterns: allowPatterns, + deny_patterns: denyPatterns, + command: testCommand, + }), + }) + const data = await res.json() + setTestResult({ + allowed: data.allowed, + blocked: data.blocked, + matchedWhitelist: data.matched_whitelist ?? null, + matchedBlacklist: data.matched_blacklist ?? null, + }) + } catch { + setTestResult(null) + } finally { + setIsLoading(false) + } + } return ( @@ -266,6 +328,50 @@ export function ExecSection({ form, onFieldChange }: ExecSectionProps) { /> + +
+
+ setTestCommand(e.target.value)} + onKeyDown={(e) => { + if (e.key === "Enter") { + testPatterns() + } + }} + /> + +
+ {testResult && ( +
+ {testResult.allowed + ? `${t("pages.config.pattern_detector_result_allowed")}${testResult.matchedWhitelist ? ` (${testResult.matchedWhitelist})` : ""}` + : testResult.blocked + ? `${t("pages.config.pattern_detector_result_blocked")}${testResult.matchedBlacklist ? ` (${testResult.matchedBlacklist})` : ""}` + : t("pages.config.pattern_detector_result_no_match")} +
+ )} +
+ + export interface CoreConfigForm { workspace: string restrictToWorkspace: boolean + splitOnMarker: boolean toolFeedbackEnabled: boolean toolFeedbackMaxArgsLength: string execEnabled: boolean @@ -65,7 +66,8 @@ export const DM_SCOPE_OPTIONS = [ export const EMPTY_FORM: CoreConfigForm = { workspace: "", restrictToWorkspace: true, - toolFeedbackEnabled: true, + splitOnMarker: false, + toolFeedbackEnabled: false, toolFeedbackMaxArgsLength: "300", execEnabled: true, allowRemote: true, @@ -136,6 +138,10 @@ export function buildFormFromConfig(config: unknown): CoreConfigForm { defaults.restrict_to_workspace === undefined ? EMPTY_FORM.restrictToWorkspace : asBool(defaults.restrict_to_workspace), + splitOnMarker: + defaults.split_on_marker === undefined + ? EMPTY_FORM.splitOnMarker + : asBool(defaults.split_on_marker), toolFeedbackEnabled: toolFeedback.enabled === undefined ? EMPTY_FORM.toolFeedbackEnabled diff --git a/web/frontend/src/components/models/add-model-sheet.tsx b/web/frontend/src/components/models/add-model-sheet.tsx index c760bc672..de9481391 100644 --- a/web/frontend/src/components/models/add-model-sheet.tsx +++ b/web/frontend/src/components/models/add-model-sheet.tsx @@ -20,6 +20,7 @@ import { SheetHeader, SheetTitle, } from "@/components/ui/sheet" +import { Textarea } from "@/components/ui/textarea" interface AddForm { modelName: string @@ -34,6 +35,7 @@ interface AddForm { maxTokensField: string requestTimeout: string thinkingLevel: string + extraBody: string } const EMPTY_ADD_FORM: AddForm = { @@ -49,6 +51,7 @@ const EMPTY_ADD_FORM: AddForm = { maxTokensField: "", requestTimeout: "", thinkingLevel: "", + extraBody: "", } interface AddModelSheetProps { @@ -100,7 +103,8 @@ export function AddModelSheet({ } const setField = - (key: keyof AddForm) => (e: React.ChangeEvent) => { + (key: keyof AddForm) => + (e: React.ChangeEvent) => { setForm((f) => ({ ...f, [key]: e.target.value })) if (fieldErrors[key]) { setFieldErrors((prev) => ({ ...prev, [key]: undefined })) @@ -129,6 +133,9 @@ export function AddModelSheet({ ? Number(form.requestTimeout) : undefined, thinking_level: form.thinkingLevel.trim() || undefined, + extra_body: form.extraBody.trim() + ? JSON.parse(form.extraBody.trim()) + : undefined, }) if (setAsDefault) { await setDefaultModel(modelName) @@ -305,6 +312,18 @@ export function AddModelSheet({ placeholder="max_completion_tokens" /> + + +