Merge branch 'sipeed:main' into main
This commit is contained in:
commit
41f3455b94
91 changed files with 8217 additions and 7590 deletions
5
.github/workflows/pr.yml
vendored
5
.github/workflows/pr.yml
vendored
|
|
@ -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 ./...
|
||||
|
|
|
|||
1
.gitignore
vendored
1
.gitignore
vendored
|
|
@ -40,6 +40,7 @@ tasks/
|
|||
|
||||
# Plans
|
||||
docs/plans/
|
||||
docs/superpowers/
|
||||
|
||||
# Editors
|
||||
.vscode/
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ builds:
|
|||
env:
|
||||
- CGO_ENABLED=0
|
||||
tags:
|
||||
- goolm
|
||||
- stdjson
|
||||
ldflags:
|
||||
- -s -w
|
||||
|
|
@ -57,6 +58,7 @@ builds:
|
|||
env:
|
||||
- CGO_ENABLED=0
|
||||
tags:
|
||||
- goolm
|
||||
- stdjson
|
||||
ldflags:
|
||||
- -s -w
|
||||
|
|
@ -95,6 +97,7 @@ builds:
|
|||
env:
|
||||
- CGO_ENABLED=0
|
||||
tags:
|
||||
- goolm
|
||||
- stdjson
|
||||
ldflags:
|
||||
- -s -w
|
||||
|
|
|
|||
62
Makefile
62
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).
|
||||
#
|
||||
|
|
@ -130,15 +136,15 @@ 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 +153,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 +179,18 @@ 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)
|
||||
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
|
||||
|
|
@ -224,13 +230,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
|
||||
|
|
@ -239,11 +245,11 @@ fmt:
|
|||
|
||||
## lint: Run linters
|
||||
lint:
|
||||
@$(GOLANGCI_LINT) run
|
||||
@CGO_ENABLED=0 $(GOLANGCI_LINT) run --build-tags $(GO_BUILD_TAGS)
|
||||
|
||||
## fix: Fix linting issues
|
||||
fix:
|
||||
@$(GOLANGCI_LINT) run --fix
|
||||
@CGO_ENABLED=0 $(GOLANGCI_LINT) run --fix --build-tags $(GO_BUILD_TAGS)
|
||||
|
||||
## deps: Download dependencies
|
||||
deps:
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ func NewAuthCommand() *cobra.Command {
|
|||
newStatusCommand(),
|
||||
newModelsCommand(),
|
||||
newWeixinCommand(),
|
||||
newWeComCommand(),
|
||||
)
|
||||
|
||||
return cmd
|
||||
|
|
|
|||
|
|
@ -33,6 +33,7 @@ func TestNewAuthCommand(t *testing.T) {
|
|||
"status",
|
||||
"models",
|
||||
"weixin",
|
||||
"wecom",
|
||||
}
|
||||
|
||||
subcommands := cmd.Commands()
|
||||
|
|
|
|||
407
cmd/picoclaw/internal/auth/wecom.go
Normal file
407
cmd/picoclaw/internal/auth/wecom.go
Normal file
|
|
@ -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
|
||||
}
|
||||
}
|
||||
157
cmd/picoclaw/internal/auth/wecom_test.go
Normal file
157
cmd/picoclaw/internal/auth/wecom_test.go
Normal file
|
|
@ -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.")
|
||||
}
|
||||
|
|
@ -162,7 +162,9 @@
|
|||
"enabled": true,
|
||||
"text": "Thinking... 💭"
|
||||
},
|
||||
"reasoning_channel_id": ""
|
||||
"reasoning_channel_id": "",
|
||||
"crypto_database_path": "",
|
||||
"crypto_passphrase": "YOUR_MATRIX_CRYPTO_PICKLE_KEY"
|
||||
},
|
||||
"line": {
|
||||
"enabled": false,
|
||||
|
|
@ -182,39 +184,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": {
|
||||
|
|
|
|||
|
|
@ -25,7 +25,9 @@ Add this to `config.json`:
|
|||
"text": "Thinking..."
|
||||
},
|
||||
"reasoning_channel_id": "",
|
||||
"message_format": "richtext"
|
||||
"message_format": "richtext",
|
||||
"crypto_database_path": "",
|
||||
"crypto_passphrase": "YOUR_MATRIX_CRYPTO_PICKLE_KEY"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -46,6 +48,8 @@ Add this to `config.json`:
|
|||
| placeholder | object | No | Placeholder message config |
|
||||
| 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 |
|
||||
|
||||
## 3. Currently Supported
|
||||
|
||||
|
|
@ -58,6 +62,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
|
||||
|
||||
|
|
|
|||
|
|
@ -24,7 +24,10 @@
|
|||
"enabled": true,
|
||||
"text": "Thinking... 💭"
|
||||
},
|
||||
"reasoning_channel_id": ""
|
||||
"reasoning_channel_id": "",
|
||||
"message_format": "richtext",
|
||||
"crypto_database_path": "",
|
||||
"crypto_passphrase": "YOUR_MATRIX_CRYPTO_PICKLE_KEY"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -45,6 +48,8 @@
|
|||
| placeholder | object | 否 | 占位消息配置 |
|
||||
| reasoning_channel_id | string | 否 | 思维链输出目标通道 |
|
||||
| message_format | string | 否 | 消息格式:`richtext`(富文本)或 `plain`(纯文本) |
|
||||
| crypto_database_path | string | 否 | 加密数据库存储路径(为空时使用工作空间路径 `~/.picoclaw/workspace`) |
|
||||
| crypto_passphrase | string | 否 | 加密数据库中 session key 的序列化密钥;设置后不能更改 |
|
||||
|
||||
## 3. 当前支持
|
||||
|
||||
|
|
@ -56,6 +61,7 @@
|
|||
- Typing 状态(`m.typing`)
|
||||
- 占位消息(`Thinking... 💭`)+ 最终回复替换
|
||||
- 自动加入邀请房间(可关闭)
|
||||
- 端对端加密(E2EE)消息支持
|
||||
|
||||
## 4. TODO
|
||||
|
||||
|
|
|
|||
3
go.mod
3
go.mod
|
|
@ -20,6 +20,7 @@ require (
|
|||
github.com/gorilla/websocket v1.5.3
|
||||
github.com/h2non/filetype v1.1.3
|
||||
github.com/larksuite/oapi-sdk-go/v3 v3.5.3
|
||||
github.com/mattn/go-sqlite3 v1.14.34
|
||||
github.com/mdp/qrterminal/v3 v3.2.1
|
||||
github.com/modelcontextprotocol/go-sdk v1.4.1
|
||||
github.com/mymmrac/telego v1.7.0
|
||||
|
|
@ -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
|
||||
|
|
@ -77,7 +79,6 @@ 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
|
||||
|
|
|
|||
|
|
@ -96,14 +96,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 +1031,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.
|
||||
|
|
@ -2165,6 +2166,7 @@ turnLoop:
|
|||
"iteration": iteration,
|
||||
})
|
||||
|
||||
allResponsesHandled := len(normalizedToolCalls) > 0
|
||||
assistantMsg := providers.Message{
|
||||
Role: "assistant",
|
||||
Content: response.Content,
|
||||
|
|
@ -2221,6 +2223,7 @@ turnLoop:
|
|||
toolArgs = toolReq.Arguments
|
||||
}
|
||||
case HookActionDenyTool:
|
||||
allResponsesHandled = false
|
||||
denyContent := hookDeniedToolContent("Tool execution denied by hook", decision.Reason)
|
||||
al.emitEvent(
|
||||
EventKindToolExecSkipped,
|
||||
|
|
@ -2260,6 +2263,7 @@ turnLoop:
|
|||
ChatID: ts.chatID,
|
||||
})
|
||||
if !approval.Approved {
|
||||
allResponsesHandled = false
|
||||
denyContent := hookDeniedToolContent("Tool execution denied by approval hook", approval.Reason)
|
||||
al.emitEvent(
|
||||
EventKindToolExecSkipped,
|
||||
|
|
@ -2333,10 +2337,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 +2421,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 +2479,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 +2574,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 +3245,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 +3376,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 +3431,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."
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -33,6 +33,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 +324,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 +570,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 +926,98 @@ 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 toolLimitOnlyProvider struct{}
|
||||
|
||||
func (m *toolLimitOnlyProvider) Chat(
|
||||
|
|
@ -636,8 +1054,9 @@ func (m *mockCustomTool) Description() string {
|
|||
|
||||
func (m *mockCustomTool) Parameters() map[string]any {
|
||||
return map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{},
|
||||
"type": "object",
|
||||
"properties": map[string]any{},
|
||||
"additionalProperties": true,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -645,6 +1064,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 {
|
||||
|
|
@ -1494,18 +2042,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})
|
||||
}
|
||||
|
|
@ -1525,7 +2072,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: ""},
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
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.
|
||||
|
|
|
|||
|
|
@ -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()` 方法暴露此配置。
|
||||
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()` 方法暴露此配置。
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
}
|
||||
|
|
@ -774,7 +799,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 +807,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 +856,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 +865,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 +876,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 +1069,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]
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
})
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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})
|
||||
|
|
@ -548,9 +651,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 +762,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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
@ -208,23 +300,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 +335,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 +360,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 +428,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
|
||||
|
|
|
|||
144
pkg/channels/pico/pico_test.go
Normal file
144
pkg/channels/pico/pico_test.go
Normal file
|
|
@ -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
|
||||
}
|
||||
|
|
@ -642,8 +642,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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
}
|
||||
}
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -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)
|
||||
}
|
||||
File diff suppressed because it is too large
Load diff
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
File diff suppressed because it is too large
Load diff
|
|
@ -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)
|
||||
}
|
||||
|
|
@ -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 := "<xml><Content>Hello</Content></xml>"
|
||||
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")
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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")
|
||||
}
|
||||
}
|
||||
|
|
@ -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)
|
||||
})
|
||||
}
|
||||
|
|
|
|||
802
pkg/channels/wecom/media.go
Normal file
802
pkg/channels/wecom/media.go
Normal file
|
|
@ -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
|
||||
}
|
||||
180
pkg/channels/wecom/media_test.go
Normal file
180
pkg/channels/wecom/media_test.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
173
pkg/channels/wecom/protocol.go
Normal file
173
pkg/channels/wecom/protocol.go
Normal file
|
|
@ -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
|
||||
}
|
||||
113
pkg/channels/wecom/reqid_store.go
Normal file
113
pkg/channels/wecom/reqid_store.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
24
pkg/channels/wecom/reqid_store_test.go
Normal file
24
pkg/channels/wecom/reqid_store_test.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
970
pkg/channels/wecom/wecom.go
Normal file
970
pkg/channels/wecom/wecom.go
Normal file
|
|
@ -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
|
||||
}
|
||||
660
pkg/channels/wecom/wecom_test.go
Normal file
660
pkg/channels/wecom/wecom_test.go
Normal file
|
|
@ -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,
|
||||
}
|
||||
}
|
||||
|
|
@ -99,13 +99,7 @@ Examples:
|
|||
- `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.wecom.secret`
|
||||
- `ref:channels.pico.token`
|
||||
- `ref:channels.irc.password`
|
||||
- `ref:channels.irc.nickserv_password`
|
||||
|
|
|
|||
|
|
@ -106,6 +106,7 @@ func (c *Config) WithSecurity(sec *SecurityConfig) *Config {
|
|||
c.security = sec
|
||||
return c
|
||||
}
|
||||
sec = normalizeSecurityConfig(sec)
|
||||
err := applySecurityConfig(c, sec)
|
||||
if err != nil {
|
||||
return nil
|
||||
|
|
@ -320,10 +321,7 @@ type AgentDefaults struct {
|
|||
ToolFeedback ToolFeedbackConfig `json:"tool_feedback,omitempty"`
|
||||
}
|
||||
|
||||
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 {
|
||||
|
|
@ -363,9 +361,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"`
|
||||
|
|
@ -590,18 +586,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
|
||||
|
|
@ -677,136 +675,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
|
||||
}
|
||||
|
|
@ -814,6 +704,7 @@ func (c *WeComAIBotConfig) SetSecret(secret string) {
|
|||
type WeixinConfig struct {
|
||||
Enabled bool `json:"enabled" env:"PICOCLAW_CHANNELS_WEIXIN_ENABLED"`
|
||||
token string
|
||||
AccountID string `json:"account_id,omitempty" env:"PICOCLAW_CHANNELS_WEIXIN_ACCOUNT_ID"`
|
||||
BaseURL string `json:"base_url" env:"PICOCLAW_CHANNELS_WEIXIN_BASE_URL"`
|
||||
CDNBaseURL string `json:"cdn_base_url" env:"PICOCLAW_CHANNELS_WEIXIN_CDN_BASE_URL"`
|
||||
Proxy string `json:"proxy" env:"PICOCLAW_CHANNELS_WEIXIN_PROXY"`
|
||||
|
|
@ -1397,6 +1288,18 @@ func LoadConfig(path string) (*Config, error) {
|
|||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Load existing security config and merge with migrated one to prevent data loss
|
||||
existingSec, secErr := loadSecurityConfig(securityPath(path))
|
||||
if secErr != nil {
|
||||
logger.WarnF("failed to load existing security config during migration", map[string]any{"error": secErr})
|
||||
}
|
||||
if existingSec != nil && cfg.security != nil {
|
||||
cfg.security = mergeSecurityConfig(existingSec, cfg.security)
|
||||
// Re-apply the merged security config to update all channels and models
|
||||
if err = applySecurityConfig(cfg, cfg.security); err != nil {
|
||||
logger.WarnF("failed to re-apply merged security config during migration", map[string]any{"error": err})
|
||||
}
|
||||
}
|
||||
defer func(cfg *Config) {
|
||||
_ = SaveConfig(path, cfg)
|
||||
}(cfg)
|
||||
|
|
@ -1613,39 +1516,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
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -1772,6 +1646,7 @@ func SaveConfig(path string, cfg *Config) error {
|
|||
logger.ErrorC("config", "security is nil")
|
||||
return fmt.Errorf("security is nil")
|
||||
}
|
||||
cfg.security = normalizeSecurityConfig(cfg.security)
|
||||
// Ensure version is always set when saving
|
||||
if cfg.Version == 0 {
|
||||
cfg.Version = CurrentVersion
|
||||
|
|
@ -1868,27 +1743,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(),
|
||||
|
|
@ -2021,6 +1879,12 @@ func (c *Config) SecurityCopyFrom(cfg *Config) {
|
|||
}
|
||||
}
|
||||
|
||||
// ApplySecurity re-applies the stored security config to populate private fields (tokens, API keys, etc.).
|
||||
// Call this after SecurityCopyFrom when you need private fields to be accessible for validation or use.
|
||||
func (c *Config) ApplySecurity() error {
|
||||
return applySecurityConfig(c, c.security)
|
||||
}
|
||||
|
||||
func MergeAPIKeys(apiKey string, apiKeys []string) []string {
|
||||
seen := make(map[string]struct{})
|
||||
var all []string
|
||||
|
|
|
|||
|
|
@ -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"`
|
||||
|
|
|
|||
|
|
@ -1372,8 +1372,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",
|
||||
|
|
|
|||
|
|
@ -113,6 +113,8 @@ func DefaultConfig() *Config {
|
|||
Enabled: true,
|
||||
Text: "Thinking... 💭",
|
||||
},
|
||||
CryptoDatabasePath: "",
|
||||
CryptoPassphrase: "",
|
||||
},
|
||||
LINE: LINEConfig{
|
||||
Enabled: false,
|
||||
|
|
@ -129,32 +131,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,
|
||||
|
|
|
|||
|
|
@ -153,13 +153,7 @@ Both single and multiple keys should use the array format.
|
|||
- 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.wecom.secret
|
||||
- ref:channels.pico.token
|
||||
- ref:channels.irc.password
|
||||
- ref:channels.irc.nickserv_password
|
||||
|
|
|
|||
|
|
@ -566,3 +566,118 @@ func TestMigration_Integration_ModelNameField(t *testing.T) {
|
|||
t.Errorf("ModelFallbacks[0] = %q, want %q", cfg.Agents.Defaults.ModelFallbacks[0], "deepseek-chat")
|
||||
}
|
||||
}
|
||||
|
||||
// TestMigration_PreservesExistingSecurityConfig tests that when migrating from v0 to v1,
|
||||
// existing .security.yml values (e.g., loaded from environment variables) are preserved
|
||||
// and not overwritten by empty values from the legacy config.
|
||||
func TestMigration_PreservesExistingSecurityConfig(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
configPath := filepath.Join(tmpDir, "config.json")
|
||||
securityPath := filepath.Join(tmpDir, ".security.yml")
|
||||
|
||||
// Create a legacy config (version 0) with model_list and channel config
|
||||
// The model_list doesn't have api_keys, they should come from existing .security.yml
|
||||
legacyConfig := `{
|
||||
"agents": {
|
||||
"defaults": {
|
||||
"provider": "openai",
|
||||
"model": "gpt-4"
|
||||
}
|
||||
},
|
||||
"model_list": [
|
||||
{
|
||||
"model_name": "openai",
|
||||
"model": "openai/gpt-4"
|
||||
}
|
||||
],
|
||||
"channels": {
|
||||
"telegram": {
|
||||
"enabled": true
|
||||
}
|
||||
},
|
||||
"gateway": {
|
||||
"host": "127.0.0.1",
|
||||
"port": 18790
|
||||
},
|
||||
"tools": {
|
||||
"web": {"enabled": true}
|
||||
},
|
||||
"heartbeat": {
|
||||
"enabled": true,
|
||||
"interval": 30
|
||||
},
|
||||
"devices": {
|
||||
"enabled": false
|
||||
}
|
||||
}`
|
||||
|
||||
// Create an existing .security.yml with values that might come from env vars
|
||||
existingSecurity := `model_list:
|
||||
openai:0:
|
||||
api_keys:
|
||||
- sk-existing-key-from-env
|
||||
channels:
|
||||
telegram:
|
||||
token: existing-telegram-token-from-env
|
||||
discord:
|
||||
token: existing-discord-token-from-env
|
||||
web:
|
||||
brave:
|
||||
api_keys:
|
||||
- existing-brave-key
|
||||
`
|
||||
|
||||
if err := os.WriteFile(configPath, []byte(legacyConfig), 0o600); err != nil {
|
||||
t.Fatalf("Failed to write legacy config: %v", err)
|
||||
}
|
||||
|
||||
if err := os.WriteFile(securityPath, []byte(existingSecurity), 0o600); err != nil {
|
||||
t.Fatalf("Failed to write existing security config: %v", err)
|
||||
}
|
||||
|
||||
// Load the config - this should trigger migration
|
||||
cfg, err := LoadConfig(configPath)
|
||||
if err != nil {
|
||||
t.Fatalf("LoadConfig failed: %v", err)
|
||||
}
|
||||
|
||||
// Verify that the migrated config has the existing security values
|
||||
// Telegram token should be preserved
|
||||
if cfg.Channels.Telegram.Token() != "existing-telegram-token-from-env" {
|
||||
t.Errorf("Telegram token was overwritten: got %q, want %q",
|
||||
cfg.Channels.Telegram.Token(), "existing-telegram-token-from-env")
|
||||
}
|
||||
|
||||
// Discord token should be preserved (even though legacy config didn't have it)
|
||||
if cfg.Channels.Discord.Token() != "existing-discord-token-from-env" {
|
||||
t.Errorf("Discord token was overwritten: got %q, want %q",
|
||||
cfg.Channels.Discord.Token(), "existing-discord-token-from-env")
|
||||
}
|
||||
|
||||
// Model API key should be preserved
|
||||
if cfg.ModelList[0].APIKey() != "sk-existing-key-from-env" {
|
||||
t.Errorf("Model API key was overwritten: got %q, want %q",
|
||||
cfg.ModelList[0].APIKey(), "sk-existing-key-from-env")
|
||||
}
|
||||
|
||||
// Brave API key should be preserved
|
||||
if cfg.Tools.Web.Brave.APIKey() != "existing-brave-key" {
|
||||
t.Errorf("Brave API key was overwritten: got %q, want %q",
|
||||
cfg.Tools.Web.Brave.APIKey(), "existing-brave-key")
|
||||
}
|
||||
|
||||
// Reload the security config from disk to verify it wasn't corrupted
|
||||
reloadedSec, err := loadSecurityConfig(securityPath)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to reload security config: %v", err)
|
||||
}
|
||||
|
||||
if reloadedSec.Channels.Telegram == nil ||
|
||||
reloadedSec.Channels.Telegram.Token != "existing-telegram-token-from-env" {
|
||||
t.Error("Telegram token not preserved in .security.yml file")
|
||||
}
|
||||
|
||||
if reloadedSec.Channels.Discord == nil || reloadedSec.Channels.Discord.Token != "existing-discord-token-from-env" {
|
||||
t.Error("Discord token not preserved in .security.yml file")
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -25,6 +25,25 @@ const (
|
|||
SecurityConfigFile = ".security.yml"
|
||||
)
|
||||
|
||||
func normalizeSecurityConfig(sec *SecurityConfig) *SecurityConfig {
|
||||
if sec == nil {
|
||||
sec = &SecurityConfig{}
|
||||
}
|
||||
if sec.ModelList == nil {
|
||||
sec.ModelList = map[string]ModelSecurityEntry{}
|
||||
}
|
||||
if sec.Channels == nil {
|
||||
sec.Channels = &ChannelsSecurity{}
|
||||
}
|
||||
if sec.Web == nil {
|
||||
sec.Web = &WebToolsSecurity{}
|
||||
}
|
||||
if sec.Skills == nil {
|
||||
sec.Skills = &SkillsSecurity{}
|
||||
}
|
||||
return sec
|
||||
}
|
||||
|
||||
// SecurityConfig stores all sensitive data (API keys, tokens, secrets, passwords)
|
||||
// This data is loaded from security.yml and kept separate from the main config
|
||||
type SecurityConfig struct {
|
||||
|
|
@ -50,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 {
|
||||
|
|
@ -112,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 {
|
||||
|
|
@ -191,7 +195,7 @@ func loadSecurityConfig(securityPath string) (*SecurityConfig, error) {
|
|||
data, err := os.ReadFile(securityPath)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
return &SecurityConfig{}, nil
|
||||
return normalizeSecurityConfig(nil), nil
|
||||
}
|
||||
return nil, fmt.Errorf("failed to read security config: %w", err)
|
||||
}
|
||||
|
|
@ -210,7 +214,7 @@ func loadSecurityConfig(securityPath string) (*SecurityConfig, error) {
|
|||
return nil, err
|
||||
}
|
||||
|
||||
return &sec, nil
|
||||
return normalizeSecurityConfig(&sec), nil
|
||||
}
|
||||
|
||||
// saveSecurityConfig saves the security configuration to security.yml
|
||||
|
|
@ -225,6 +229,134 @@ func saveSecurityConfig(securityPath string, sec *SecurityConfig) error {
|
|||
return fileutil.WriteFileAtomic(securityPath, buf.Bytes(), 0o600)
|
||||
}
|
||||
|
||||
// mergeSecurityConfig merges two SecurityConfig instances, preferring non-empty values from 'newer'.
|
||||
// This is used during config migration to preserve existing security data while adding new entries.
|
||||
func mergeSecurityConfig(existing, newer *SecurityConfig) *SecurityConfig {
|
||||
if existing == nil {
|
||||
return normalizeSecurityConfig(newer)
|
||||
}
|
||||
if newer == nil {
|
||||
return normalizeSecurityConfig(existing)
|
||||
}
|
||||
|
||||
result := normalizeSecurityConfig(nil)
|
||||
|
||||
// Merge ModelList: prefer newer if it has keys, otherwise use existing
|
||||
for k, v := range existing.ModelList {
|
||||
result.ModelList[k] = v
|
||||
}
|
||||
for k, v := range newer.ModelList {
|
||||
if len(v.APIKeys) > 0 {
|
||||
result.ModelList[k] = v
|
||||
}
|
||||
}
|
||||
|
||||
// Merge Channels
|
||||
if existing.Channels != nil {
|
||||
result.Channels = existing.Channels
|
||||
}
|
||||
if newer.Channels != nil {
|
||||
if result.Channels == nil {
|
||||
result.Channels = &ChannelsSecurity{}
|
||||
}
|
||||
mergeChannelsSecurity(result.Channels, newer.Channels)
|
||||
}
|
||||
|
||||
// Merge Web
|
||||
if existing.Web != nil {
|
||||
result.Web = existing.Web
|
||||
}
|
||||
if newer.Web != nil {
|
||||
if result.Web == nil {
|
||||
result.Web = &WebToolsSecurity{}
|
||||
}
|
||||
mergeWebToolsSecurity(result.Web, newer.Web)
|
||||
}
|
||||
|
||||
// Merge Skills
|
||||
if existing.Skills != nil {
|
||||
result.Skills = existing.Skills
|
||||
}
|
||||
if newer.Skills != nil {
|
||||
if result.Skills == nil {
|
||||
result.Skills = &SkillsSecurity{}
|
||||
}
|
||||
mergeSkillsSecurity(result.Skills, newer.Skills)
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
func mergeChannelsSecurity(dst, src *ChannelsSecurity) {
|
||||
if src.Telegram != nil && src.Telegram.Token != "" {
|
||||
dst.Telegram = src.Telegram
|
||||
}
|
||||
if src.Feishu != nil &&
|
||||
(src.Feishu.AppSecret != "" || src.Feishu.EncryptKey != "" || src.Feishu.VerificationToken != "") {
|
||||
dst.Feishu = src.Feishu
|
||||
}
|
||||
if src.Discord != nil && src.Discord.Token != "" {
|
||||
dst.Discord = src.Discord
|
||||
}
|
||||
if src.Weixin != nil && src.Weixin.Token != "" {
|
||||
dst.Weixin = src.Weixin
|
||||
}
|
||||
if src.QQ != nil && src.QQ.AppSecret != "" {
|
||||
dst.QQ = src.QQ
|
||||
}
|
||||
if src.DingTalk != nil && src.DingTalk.ClientSecret != "" {
|
||||
dst.DingTalk = src.DingTalk
|
||||
}
|
||||
if src.Slack != nil && (src.Slack.BotToken != "" || src.Slack.AppToken != "") {
|
||||
dst.Slack = src.Slack
|
||||
}
|
||||
if src.Matrix != nil && src.Matrix.AccessToken != "" {
|
||||
dst.Matrix = src.Matrix
|
||||
}
|
||||
if src.LINE != nil && (src.LINE.ChannelSecret != "" || src.LINE.ChannelAccessToken != "") {
|
||||
dst.LINE = src.LINE
|
||||
}
|
||||
if src.OneBot != nil && src.OneBot.AccessToken != "" {
|
||||
dst.OneBot = src.OneBot
|
||||
}
|
||||
if src.WeCom != nil && src.WeCom.Secret != "" {
|
||||
dst.WeCom = src.WeCom
|
||||
}
|
||||
if src.Pico != nil && src.Pico.Token != "" {
|
||||
dst.Pico = src.Pico
|
||||
}
|
||||
if src.IRC != nil && (src.IRC.Password != "" || src.IRC.NickServPassword != "" || src.IRC.SASLPassword != "") {
|
||||
dst.IRC = src.IRC
|
||||
}
|
||||
}
|
||||
|
||||
func mergeWebToolsSecurity(dst, src *WebToolsSecurity) {
|
||||
if src.Brave != nil && len(src.Brave.APIKeys) > 0 {
|
||||
dst.Brave = src.Brave
|
||||
}
|
||||
if src.Tavily != nil && len(src.Tavily.APIKeys) > 0 {
|
||||
dst.Tavily = src.Tavily
|
||||
}
|
||||
if src.Perplexity != nil && len(src.Perplexity.APIKeys) > 0 {
|
||||
dst.Perplexity = src.Perplexity
|
||||
}
|
||||
if src.GLMSearch != nil && src.GLMSearch.APIKey != "" {
|
||||
dst.GLMSearch = src.GLMSearch
|
||||
}
|
||||
if src.BaiduSearch != nil && src.BaiduSearch.APIKey != "" {
|
||||
dst.BaiduSearch = src.BaiduSearch
|
||||
}
|
||||
}
|
||||
|
||||
func mergeSkillsSecurity(dst, src *SkillsSecurity) {
|
||||
if src.Github != nil && src.Github.Token != "" {
|
||||
dst.Github = src.Github
|
||||
}
|
||||
if src.ClawHub != nil && src.ClawHub.AuthToken != "" {
|
||||
dst.ClawHub = src.ClawHub
|
||||
}
|
||||
}
|
||||
|
||||
// SensitiveDataCache caches the compiled regex for filtering sensitive data.
|
||||
// SensitiveDataCache caches the strings.Replacer for filtering sensitive data.
|
||||
// Computed once on first access via sync.Once.
|
||||
|
|
|
|||
|
|
@ -17,13 +17,12 @@ import (
|
|||
|
||||
// Test JSON unmarshal of private fields
|
||||
func TestJSONUnmarshalPrivateFields(t *testing.T) {
|
||||
//nolint: govet
|
||||
type testStruct struct {
|
||||
PublicField string `json:"public"`
|
||||
privateField string `json:"private"`
|
||||
privateField string
|
||||
}
|
||||
|
||||
data := `{"public": "pub", "private": "priv"}`
|
||||
data := `{"public": "pub", "privateField": "priv"}`
|
||||
var s testStruct
|
||||
if err := json.Unmarshal([]byte(data), &s); err != nil {
|
||||
t.Fatalf("JSON unmarshal failed: %v", err)
|
||||
|
|
@ -35,9 +34,8 @@ func TestJSONUnmarshalPrivateFields(t *testing.T) {
|
|||
if s.PublicField != "pub" {
|
||||
t.Errorf("PublicField = %q, want 'pub'", s.PublicField)
|
||||
}
|
||||
// This should fail because privateField is unexported
|
||||
if s.privateField != "priv" {
|
||||
t.Logf("privateField = %q, want 'priv' - THIS IS EXPECTED TO FAIL", s.privateField)
|
||||
if s.privateField != "" {
|
||||
t.Errorf("privateField = %q, want empty because unexported fields are ignored", s.privateField)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -242,15 +240,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
|
||||
|
|
@ -317,15 +307,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:
|
||||
|
|
@ -411,24 +393,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())
|
||||
|
|
|
|||
|
|
@ -20,6 +20,9 @@ func TestSecurityConfig(t *testing.T) {
|
|||
require.NoError(t, err)
|
||||
assert.NotNil(t, sec)
|
||||
assert.Empty(t, sec.ModelList)
|
||||
assert.NotNil(t, sec.Channels)
|
||||
assert.NotNil(t, sec.Web)
|
||||
assert.NotNil(t, sec.Skills)
|
||||
})
|
||||
}
|
||||
|
||||
|
|
|
|||
21
pkg/gateway/channel_matrix.go
Normal file
21
pkg/gateway/channel_matrix.go
Normal file
|
|
@ -0,0 +1,21 @@
|
|||
//go:build !mipsle && !netbsd
|
||||
|
||||
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).
|
||||
//
|
||||
// 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"
|
||||
)
|
||||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
292
pkg/tools/normalization.go
Normal file
292
pkg/tools/normalization.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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) {
|
||||
|
|
@ -180,6 +206,14 @@ func (r *ToolRegistry) ExecuteWithContext(
|
|||
return ErrorResult(fmt.Sprintf("tool %q not found", name)).WithError(fmt.Errorf("tool not found"))
|
||||
}
|
||||
|
||||
// Validate arguments against the tool's declared schema.
|
||||
if err := validateToolArgs(tool.Parameters(), args); err != nil {
|
||||
logger.WarnCF("tool", "Tool argument validation failed",
|
||||
map[string]any{"tool": name, "error": err.Error()})
|
||||
return ErrorResult(fmt.Sprintf("invalid arguments for tool %q: %s", name, err)).
|
||||
WithError(fmt.Errorf("argument validation failed: %w", err))
|
||||
}
|
||||
|
||||
// Inject channel/chatID into ctx so tools read them via ToolChannel(ctx)/ToolChatID(ctx).
|
||||
// Always inject — tools validate what they require.
|
||||
ctx = WithToolContext(ctx, channel, chatID)
|
||||
|
|
@ -230,6 +264,8 @@ func (r *ToolRegistry) ExecuteWithContext(
|
|||
}
|
||||
}
|
||||
|
||||
result = normalizeToolResult(result, name, r.mediaStore, channel, chatID)
|
||||
|
||||
duration := time.Since(start)
|
||||
|
||||
// Log based on result type
|
||||
|
|
@ -251,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()),
|
||||
})
|
||||
}
|
||||
|
||||
|
|
@ -346,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{
|
||||
|
|
|
|||
|
|
@ -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 := ""
|
||||
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  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)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
209
pkg/tools/validate.go
Normal file
209
pkg/tools/validate.go
Normal file
|
|
@ -0,0 +1,209 @@
|
|||
package tools
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
)
|
||||
|
||||
// validateToolArgs validates args against a JSON Schema-like map.
|
||||
// schema is expected to have optional keys: "properties", "required", "additionalProperties".
|
||||
func validateToolArgs(schema map[string]any, args map[string]any) error {
|
||||
if len(schema) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
if args == nil {
|
||||
args = map[string]any{}
|
||||
}
|
||||
|
||||
if err := checkRequired(schema, args); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
propsRaw, ok := schema["properties"]
|
||||
if !ok {
|
||||
return nil // no properties defined — accept any args
|
||||
}
|
||||
|
||||
props, ok := propsRaw.(map[string]any)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
|
||||
additional := allowsAdditional(schema)
|
||||
|
||||
for key, val := range args {
|
||||
propSchemaRaw, known := props[key]
|
||||
if !known {
|
||||
if !additional {
|
||||
return fmt.Errorf("unexpected property %q", key)
|
||||
}
|
||||
continue
|
||||
}
|
||||
propSchema, ok := propSchemaRaw.(map[string]any)
|
||||
if !ok {
|
||||
continue // can't validate without a proper schema map
|
||||
}
|
||||
if err := checkType(key, val, propSchema); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// checkRequired verifies that every field listed in schema["required"] is present in args.
|
||||
func checkRequired(schema map[string]any, args map[string]any) error {
|
||||
reqRaw, ok := schema["required"]
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
|
||||
var required []string
|
||||
|
||||
switch r := reqRaw.(type) {
|
||||
case []string:
|
||||
required = r
|
||||
case []any:
|
||||
for _, v := range r {
|
||||
s, ok := v.(string)
|
||||
if ok {
|
||||
required = append(required, s)
|
||||
}
|
||||
}
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
|
||||
for _, field := range required {
|
||||
if _, present := args[field]; !present {
|
||||
return fmt.Errorf("missing required property %q", field)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// allowsAdditional returns true when the schema explicitly sets
|
||||
// "additionalProperties" to true, or when the key is absent (default: reject extras).
|
||||
func allowsAdditional(schema map[string]any) bool {
|
||||
v, ok := schema["additionalProperties"]
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
b, ok := v.(bool)
|
||||
return ok && b
|
||||
}
|
||||
|
||||
// checkType validates that val matches the JSON Schema type declared in propSchema.
|
||||
func checkType(key string, val any, propSchema map[string]any) error {
|
||||
typeRaw, ok := propSchema["type"]
|
||||
if !ok {
|
||||
return nil // no type constraint
|
||||
}
|
||||
typeName, ok := typeRaw.(string)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
|
||||
switch typeName {
|
||||
case "string":
|
||||
if _, ok := val.(string); !ok {
|
||||
return fmt.Errorf("property %q: expected string, got %T", key, val)
|
||||
}
|
||||
case "integer":
|
||||
switch v := val.(type) {
|
||||
case float64:
|
||||
if v != math.Trunc(v) {
|
||||
return fmt.Errorf("property %q: expected integer, got float64 with fractional part", key)
|
||||
}
|
||||
case int:
|
||||
// ok
|
||||
case int64:
|
||||
// ok
|
||||
default:
|
||||
return fmt.Errorf("property %q: expected integer, got %T", key, val)
|
||||
}
|
||||
case "number":
|
||||
switch val.(type) {
|
||||
case float64, int, int64:
|
||||
// ok
|
||||
default:
|
||||
return fmt.Errorf("property %q: expected number, got %T", key, val)
|
||||
}
|
||||
case "boolean":
|
||||
if _, ok := val.(bool); !ok {
|
||||
return fmt.Errorf("property %q: expected boolean, got %T", key, val)
|
||||
}
|
||||
case "array":
|
||||
arr, ok := val.([]any)
|
||||
if !ok {
|
||||
return fmt.Errorf("property %q: expected array, got %T", key, val)
|
||||
}
|
||||
if err := checkArrayItems(key, arr, propSchema); err != nil {
|
||||
return err
|
||||
}
|
||||
case "object":
|
||||
obj, ok := val.(map[string]any)
|
||||
if !ok {
|
||||
return fmt.Errorf("property %q: expected object, got %T", key, val)
|
||||
}
|
||||
if err := validateToolArgs(propSchema, obj); err != nil {
|
||||
return fmt.Errorf("property %q: %w", key, err)
|
||||
}
|
||||
}
|
||||
|
||||
if err := checkEnum(key, val, propSchema); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// checkArrayItems validates each element of arr against the "items" sub-schema.
|
||||
func checkArrayItems(key string, arr []any, propSchema map[string]any) error {
|
||||
itemsRaw, ok := propSchema["items"]
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
itemSchema, ok := itemsRaw.(map[string]any)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
for i, elem := range arr {
|
||||
elemKey := fmt.Sprintf("%s[%d]", key, i)
|
||||
if err := checkType(elemKey, elem, itemSchema); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// checkEnum validates that val is one of the allowed enum values in propSchema.
|
||||
func checkEnum(key string, val any, propSchema map[string]any) error {
|
||||
enumRaw, ok := propSchema["enum"]
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
|
||||
switch ev := enumRaw.(type) {
|
||||
case []any:
|
||||
for _, allowed := range ev {
|
||||
if val == allowed {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
case []string:
|
||||
s, ok := val.(string)
|
||||
if ok {
|
||||
for _, allowed := range ev {
|
||||
if s == allowed {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
default:
|
||||
return nil // unknown enum format, skip
|
||||
}
|
||||
|
||||
return fmt.Errorf("property %q: value %v is not in enum", key, val)
|
||||
}
|
||||
465
pkg/tools/validate_test.go
Normal file
465
pkg/tools/validate_test.go
Normal file
|
|
@ -0,0 +1,465 @@
|
|||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// Ensure imports are used.
|
||||
var (
|
||||
_ = context.Background
|
||||
_ = strings.Contains
|
||||
)
|
||||
|
||||
func TestValidateToolArgs(t *testing.T) {
|
||||
baseSchema := map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"name": map[string]any{"type": "string"},
|
||||
"age": map[string]any{"type": "integer"},
|
||||
},
|
||||
"required": []string{"name"},
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
schema map[string]any
|
||||
args map[string]any
|
||||
wantErr string // empty means no error expected
|
||||
}{
|
||||
{
|
||||
name: "valid args all required present",
|
||||
schema: baseSchema,
|
||||
args: map[string]any{"name": "alice", "age": float64(30)},
|
||||
},
|
||||
{
|
||||
name: "missing required field",
|
||||
schema: baseSchema,
|
||||
args: map[string]any{"age": float64(30)},
|
||||
wantErr: "missing required property \"name\"",
|
||||
},
|
||||
{
|
||||
name: "wrong type string field gets number",
|
||||
schema: baseSchema,
|
||||
args: map[string]any{"name": float64(42)},
|
||||
wantErr: "expected string",
|
||||
},
|
||||
{
|
||||
name: "nil args with required fields",
|
||||
schema: baseSchema,
|
||||
args: nil,
|
||||
wantErr: "missing required property \"name\"",
|
||||
},
|
||||
{
|
||||
name: "nil args no required fields",
|
||||
schema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"name": map[string]any{"type": "string"},
|
||||
},
|
||||
},
|
||||
args: nil,
|
||||
},
|
||||
{
|
||||
name: "empty args no required fields",
|
||||
schema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"name": map[string]any{"type": "string"},
|
||||
},
|
||||
},
|
||||
args: map[string]any{},
|
||||
},
|
||||
{
|
||||
name: "optional field correct type",
|
||||
schema: baseSchema,
|
||||
args: map[string]any{"name": "bob", "age": float64(25)},
|
||||
},
|
||||
{
|
||||
name: "optional field wrong type",
|
||||
schema: baseSchema,
|
||||
args: map[string]any{"name": "bob", "age": "twenty"},
|
||||
wantErr: "expected integer",
|
||||
},
|
||||
{
|
||||
name: "integer as float64 no fractional part",
|
||||
schema: baseSchema,
|
||||
args: map[string]any{"name": "carol", "age": float64(42)},
|
||||
},
|
||||
{
|
||||
name: "actual float for integer field",
|
||||
schema: baseSchema,
|
||||
args: map[string]any{"name": "dave", "age": float64(42.5)},
|
||||
wantErr: "expected integer, got float64 with fractional part",
|
||||
},
|
||||
{
|
||||
name: "number type accepts float",
|
||||
schema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"score": map[string]any{"type": "number"},
|
||||
},
|
||||
},
|
||||
args: map[string]any{"score": float64(3.14)},
|
||||
},
|
||||
{
|
||||
name: "number type accepts integer",
|
||||
schema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"score": map[string]any{"type": "number"},
|
||||
},
|
||||
},
|
||||
args: map[string]any{"score": float64(10)},
|
||||
},
|
||||
{
|
||||
name: "boolean type valid",
|
||||
schema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"flag": map[string]any{"type": "boolean"},
|
||||
},
|
||||
},
|
||||
args: map[string]any{"flag": true},
|
||||
},
|
||||
{
|
||||
name: "boolean type wrong",
|
||||
schema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"flag": map[string]any{"type": "boolean"},
|
||||
},
|
||||
},
|
||||
args: map[string]any{"flag": "true"},
|
||||
wantErr: "expected boolean",
|
||||
},
|
||||
{
|
||||
name: "required as []any from MCP deserialization",
|
||||
schema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"cmd": map[string]any{"type": "string"},
|
||||
},
|
||||
"required": []any{"cmd"},
|
||||
},
|
||||
args: map[string]any{},
|
||||
wantErr: "missing required property \"cmd\"",
|
||||
},
|
||||
{
|
||||
name: "enum valid value []any",
|
||||
schema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"color": map[string]any{"type": "string", "enum": []any{"red", "green", "blue"}},
|
||||
},
|
||||
},
|
||||
args: map[string]any{"color": "red"},
|
||||
},
|
||||
{
|
||||
name: "enum invalid value []any",
|
||||
schema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"color": map[string]any{"type": "string", "enum": []any{"red", "green", "blue"}},
|
||||
},
|
||||
},
|
||||
args: map[string]any{"color": "yellow"},
|
||||
wantErr: "not in enum",
|
||||
},
|
||||
{
|
||||
name: "enum valid value []string",
|
||||
schema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"color": map[string]any{"type": "string", "enum": []string{"red", "green", "blue"}},
|
||||
},
|
||||
},
|
||||
args: map[string]any{"color": "green"},
|
||||
},
|
||||
{
|
||||
name: "enum invalid value []string",
|
||||
schema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"color": map[string]any{"type": "string", "enum": []string{"red", "green", "blue"}},
|
||||
},
|
||||
},
|
||||
args: map[string]any{"color": "yellow"},
|
||||
wantErr: "not in enum",
|
||||
},
|
||||
{
|
||||
name: "extra unexpected property rejected",
|
||||
schema: baseSchema,
|
||||
args: map[string]any{"name": "eve", "hobby": "chess"},
|
||||
wantErr: "unexpected property \"hobby\"",
|
||||
},
|
||||
{
|
||||
name: "extra property allowed with additionalProperties true",
|
||||
schema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"name": map[string]any{"type": "string"},
|
||||
},
|
||||
"additionalProperties": true,
|
||||
},
|
||||
args: map[string]any{"name": "eve", "hobby": "chess"},
|
||||
},
|
||||
{
|
||||
name: "nested object valid",
|
||||
schema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"address": map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"city": map[string]any{"type": "string"},
|
||||
},
|
||||
"required": []string{"city"},
|
||||
},
|
||||
},
|
||||
},
|
||||
args: map[string]any{
|
||||
"address": map[string]any{"city": "Berlin"},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "nested object wrong type",
|
||||
schema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"address": map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"city": map[string]any{"type": "string"},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
args: map[string]any{"address": "not an object"},
|
||||
wantErr: "expected object",
|
||||
},
|
||||
{
|
||||
name: "array with valid element types",
|
||||
schema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"tags": map[string]any{
|
||||
"type": "array",
|
||||
"items": map[string]any{"type": "string"},
|
||||
},
|
||||
},
|
||||
},
|
||||
args: map[string]any{"tags": []any{"a", "b", "c"}},
|
||||
},
|
||||
{
|
||||
name: "array with wrong element types",
|
||||
schema: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"tags": map[string]any{
|
||||
"type": "array",
|
||||
"items": map[string]any{"type": "string"},
|
||||
},
|
||||
},
|
||||
},
|
||||
args: map[string]any{"tags": []any{"a", float64(2)}},
|
||||
wantErr: "expected string",
|
||||
},
|
||||
{
|
||||
name: "schema with no properties key accepts any args",
|
||||
schema: map[string]any{
|
||||
"type": "object",
|
||||
},
|
||||
args: map[string]any{"anything": "goes"},
|
||||
},
|
||||
{
|
||||
name: "empty schema accepts anything",
|
||||
schema: map[string]any{},
|
||||
args: map[string]any{"foo": "bar"},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
err := validateToolArgs(tc.schema, tc.args)
|
||||
if tc.wantErr == "" {
|
||||
if err != nil {
|
||||
t.Fatalf("expected no error, got: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
if err == nil {
|
||||
t.Fatalf("expected error containing %q, got nil", tc.wantErr)
|
||||
}
|
||||
if !strings.Contains(err.Error(), tc.wantErr) {
|
||||
t.Fatalf("expected error containing %q, got: %v", tc.wantErr, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateToolArgs_RegistryIntegration(t *testing.T) {
|
||||
r := NewToolRegistry()
|
||||
r.Register(&mockRegistryTool{
|
||||
name: "read_file",
|
||||
desc: "reads a file",
|
||||
params: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"path": map[string]any{"type": "string"},
|
||||
},
|
||||
"required": []string{"path"},
|
||||
},
|
||||
result: SilentResult("file contents"),
|
||||
})
|
||||
|
||||
// Valid args — should succeed
|
||||
result := r.Execute(context.Background(), "read_file", map[string]any{"path": "/tmp/x"})
|
||||
if result.IsError {
|
||||
t.Errorf("expected success, got error: %s", result.ForLLM)
|
||||
}
|
||||
|
||||
// Missing required field — should fail with validation error
|
||||
result = r.Execute(context.Background(), "read_file", map[string]any{})
|
||||
if !result.IsError {
|
||||
t.Error("expected validation error for missing required field")
|
||||
}
|
||||
if !strings.Contains(result.ForLLM, "missing required p") {
|
||||
t.Errorf("expected 'missing required p...' in error, got %q", result.ForLLM)
|
||||
}
|
||||
if result.Err == nil {
|
||||
t.Error("expected Err to be set via WithError")
|
||||
}
|
||||
|
||||
// Wrong type — should fail with validation error
|
||||
result = r.Execute(context.Background(), "read_file", map[string]any{"path": 123.0})
|
||||
if !result.IsError {
|
||||
t.Error("expected validation error for wrong type")
|
||||
}
|
||||
if !strings.Contains(result.ForLLM, "expected string") {
|
||||
t.Errorf("expected 'expected string' in error, got %q", result.ForLLM)
|
||||
}
|
||||
|
||||
// Extra property — should fail with validation error
|
||||
result = r.Execute(context.Background(), "read_file", map[string]any{"path": "/x", "__inject": true})
|
||||
if !result.IsError {
|
||||
t.Error("expected validation error for extra property")
|
||||
}
|
||||
if !strings.Contains(result.ForLLM, "unexpected prop") {
|
||||
t.Errorf("expected 'unexpected prop...' in error, got %q", result.ForLLM)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateToolArgs_RealSchemas(t *testing.T) {
|
||||
execSchema := map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"command": map[string]any{"type": "string"},
|
||||
"working_dir": map[string]any{"type": "string"},
|
||||
},
|
||||
"required": []string{"command"},
|
||||
}
|
||||
|
||||
cronSchema := map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"action": map[string]any{
|
||||
"type": "string",
|
||||
"enum": []any{"add", "list", "remove", "enable", "disable"},
|
||||
},
|
||||
},
|
||||
"required": []string{"action"},
|
||||
}
|
||||
|
||||
webSearchSchema := map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"query": map[string]any{"type": "string"},
|
||||
"count": map[string]any{"type": "integer"},
|
||||
},
|
||||
"required": []string{"query"},
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
schema map[string]any
|
||||
args map[string]any
|
||||
wantErr string
|
||||
}{
|
||||
// ExecTool
|
||||
{
|
||||
name: "exec valid args",
|
||||
schema: execSchema,
|
||||
args: map[string]any{"command": "ls -la", "working_dir": "/tmp"},
|
||||
},
|
||||
{
|
||||
name: "exec missing required command",
|
||||
schema: execSchema,
|
||||
args: map[string]any{"working_dir": "/tmp"},
|
||||
wantErr: "missing required property \"command\"",
|
||||
},
|
||||
{
|
||||
name: "exec wrong type for command",
|
||||
schema: execSchema,
|
||||
args: map[string]any{"command": float64(123)},
|
||||
wantErr: "expected string",
|
||||
},
|
||||
{
|
||||
name: "exec extra injected arg",
|
||||
schema: execSchema,
|
||||
args: map[string]any{"command": "ls", "malicious": "payload"},
|
||||
wantErr: "unexpected property \"malicious\"",
|
||||
},
|
||||
|
||||
// CronTool
|
||||
{
|
||||
name: "cron valid enum value",
|
||||
schema: cronSchema,
|
||||
args: map[string]any{"action": "add"},
|
||||
},
|
||||
{
|
||||
name: "cron invalid enum value",
|
||||
schema: cronSchema,
|
||||
args: map[string]any{"action": "destroy"},
|
||||
wantErr: "not in enum",
|
||||
},
|
||||
|
||||
// WebSearchTool
|
||||
{
|
||||
name: "websearch valid args",
|
||||
schema: webSearchSchema,
|
||||
args: map[string]any{"query": "golang testing", "count": float64(10)},
|
||||
},
|
||||
{
|
||||
name: "websearch missing required query",
|
||||
schema: webSearchSchema,
|
||||
args: map[string]any{"count": float64(5)},
|
||||
wantErr: "missing required property \"query\"",
|
||||
},
|
||||
{
|
||||
name: "websearch wrong type for count",
|
||||
schema: webSearchSchema,
|
||||
args: map[string]any{"query": "test", "count": "ten"},
|
||||
wantErr: "expected integer",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
err := validateToolArgs(tc.schema, tc.args)
|
||||
if tc.wantErr == "" {
|
||||
if err != nil {
|
||||
t.Fatalf("expected no error, got: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
if err == nil {
|
||||
t.Fatalf("expected error containing %q, got nil", tc.wantErr)
|
||||
}
|
||||
if !strings.Contains(err.Error(), tc.wantErr) {
|
||||
t.Fatalf("expected error containing %q, got: %v", tc.wantErr, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
@ -12,6 +12,7 @@ type channelCatalogItem struct {
|
|||
}
|
||||
|
||||
var channelCatalog = []channelCatalogItem{
|
||||
{Name: "weixin", ConfigKey: "weixin"},
|
||||
{Name: "telegram", ConfigKey: "telegram"},
|
||||
{Name: "discord", ConfigKey: "discord"},
|
||||
{Name: "slack", ConfigKey: "slack"},
|
||||
|
|
@ -21,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"},
|
||||
|
|
|
|||
|
|
@ -152,9 +152,13 @@ func (h *Handler) handlePatchConfig(w http.ResponseWriter, r *http.Request) {
|
|||
return
|
||||
}
|
||||
|
||||
// Copy security credentials before validation so security-managed
|
||||
// fields (e.g. pico token) are available for validation checks.
|
||||
// Restore security fields (tokens/keys) from the loaded config before validation,
|
||||
// because private fields are lost during JSON round-trip.
|
||||
newCfg.SecurityCopyFrom(cfg)
|
||||
if err := newCfg.ApplySecurity(); err != nil {
|
||||
http.Error(w, fmt.Sprintf("Failed to apply security config: %v", err), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
if errs := validateConfig(&newCfg); len(errs) > 0 {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
|
|
@ -205,6 +209,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(
|
||||
|
|
|
|||
|
|
@ -407,7 +407,7 @@ func (h *Handler) startGatewayLocked(initialStatus string, existingPid int) (int
|
|||
gateway.logs.Reset()
|
||||
|
||||
// Ensure Pico Channel is configured before starting gateway
|
||||
if _, err := h.ensurePicoChannel(""); err != nil {
|
||||
if _, err := h.EnsurePicoChannel(""); err != nil {
|
||||
logger.ErrorC("gateway", fmt.Sprintf("Warning: failed to ensure pico channel: %v", err))
|
||||
// Non-fatal: gateway can still start without pico channel
|
||||
}
|
||||
|
|
|
|||
|
|
@ -108,7 +108,12 @@ func (h *Handler) handleAddModel(w http.ResponseWriter, r *http.Request) {
|
|||
}
|
||||
defer r.Body.Close()
|
||||
|
||||
var mc config.ModelConfig
|
||||
type custom struct {
|
||||
config.ModelConfig
|
||||
APIKey string `json:"api_key"`
|
||||
}
|
||||
|
||||
var mc custom
|
||||
if err = json.Unmarshal(body, &mc); err != nil {
|
||||
http.Error(w, fmt.Sprintf("Invalid JSON: %v", err), http.StatusBadRequest)
|
||||
return
|
||||
|
|
@ -119,13 +124,17 @@ func (h *Handler) handleAddModel(w http.ResponseWriter, r *http.Request) {
|
|||
return
|
||||
}
|
||||
|
||||
if mc.APIKey != "" {
|
||||
mc.ModelConfig.SetAPIKey(mc.APIKey)
|
||||
}
|
||||
|
||||
cfg, err := config.LoadConfig(h.configPath)
|
||||
if err != nil {
|
||||
http.Error(w, fmt.Sprintf("Failed to load config: %v", err), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
cfg.ModelList = append(cfg.ModelList, &mc)
|
||||
cfg.ModelList = append(cfg.ModelList, &mc.ModelConfig)
|
||||
|
||||
if err := config.SaveConfig(h.configPath, cfg); err != nil {
|
||||
http.Error(w, fmt.Sprintf("Failed to save config: %v", err), http.StatusInternalServerError)
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
package api
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
|
|
@ -317,6 +318,44 @@ func TestHandleListModels_NormalizesWildcardLocalAPIBaseForProbe(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestHandleAddModel_PersistsAPIKey(t *testing.T) {
|
||||
configPath, cleanup := setupOAuthTestEnv(t)
|
||||
defer cleanup()
|
||||
|
||||
h := NewHandler(configPath)
|
||||
mux := http.NewServeMux()
|
||||
h.RegisterRoutes(mux)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/models", bytes.NewBufferString(`{
|
||||
"model_name":"new-model",
|
||||
"model":"openai/gpt-4o-mini",
|
||||
"api_key":"sk-new-model-key"
|
||||
}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
mux.ServeHTTP(rec, req)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
|
||||
cfg, err := config.LoadConfig(configPath)
|
||||
if err != nil {
|
||||
t.Fatalf("LoadConfig() error = %v", err)
|
||||
}
|
||||
if len(cfg.ModelList) != 2 {
|
||||
t.Fatalf("len(model_list) = %d, want 2", len(cfg.ModelList))
|
||||
}
|
||||
|
||||
added := cfg.ModelList[1]
|
||||
if added.ModelName != "new-model" {
|
||||
t.Fatalf("model_name = %q, want %q", added.ModelName, "new-model")
|
||||
}
|
||||
if added.APIKey() != "sk-new-model-key" {
|
||||
t.Fatalf("api_key = %q, want %q", added.APIKey(), "sk-new-model-key")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMaskAPIKey(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
|
|
|
|||
|
|
@ -90,14 +90,14 @@ func (h *Handler) handleRegenPicoToken(w http.ResponseWriter, r *http.Request) {
|
|||
})
|
||||
}
|
||||
|
||||
// ensurePicoChannel enables the Pico channel with sane defaults if it isn't
|
||||
// EnsurePicoChannel enables the Pico channel with sane defaults if it isn't
|
||||
// already configured. Returns true when the config was modified.
|
||||
//
|
||||
// callerOrigin is the Origin header from the setup request. If non-empty and
|
||||
// no origins are configured yet, it's written as the allowed origin so the
|
||||
// WebSocket handshake works for whatever host the caller is on (LAN, custom
|
||||
// port, etc.). Pass "" when there's no request context.
|
||||
func (h *Handler) ensurePicoChannel(callerOrigin string) (bool, error) {
|
||||
func (h *Handler) EnsurePicoChannel(callerOrigin string) (bool, error) {
|
||||
cfg, err := config.LoadConfig(h.configPath)
|
||||
if err != nil {
|
||||
return false, fmt.Errorf("failed to load config: %w", err)
|
||||
|
|
@ -134,7 +134,7 @@ func (h *Handler) ensurePicoChannel(callerOrigin string) (bool, error) {
|
|||
//
|
||||
// POST /api/pico/setup
|
||||
func (h *Handler) handlePicoSetup(w http.ResponseWriter, r *http.Request) {
|
||||
changed, err := h.ensurePicoChannel(r.Header.Get("Origin"))
|
||||
changed, err := h.EnsurePicoChannel(r.Header.Get("Origin"))
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ import (
|
|||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"testing"
|
||||
|
|
@ -17,12 +18,12 @@ func TestEnsurePicoChannel_FreshConfig(t *testing.T) {
|
|||
configPath := filepath.Join(t.TempDir(), "config.json")
|
||||
h := NewHandler(configPath)
|
||||
|
||||
changed, err := h.ensurePicoChannel("")
|
||||
changed, err := h.EnsurePicoChannel("")
|
||||
if err != nil {
|
||||
t.Fatalf("ensurePicoChannel() error = %v", err)
|
||||
t.Fatalf("EnsurePicoChannel() error = %v", err)
|
||||
}
|
||||
if !changed {
|
||||
t.Fatal("ensurePicoChannel() should report changed on a fresh config")
|
||||
t.Fatal("EnsurePicoChannel() should report changed on a fresh config")
|
||||
}
|
||||
|
||||
cfg, err := config.LoadConfig(configPath)
|
||||
|
|
@ -42,8 +43,8 @@ func TestEnsurePicoChannel_DoesNotEnableTokenQuery(t *testing.T) {
|
|||
configPath := filepath.Join(t.TempDir(), "config.json")
|
||||
h := NewHandler(configPath)
|
||||
|
||||
if _, err := h.ensurePicoChannel(""); err != nil {
|
||||
t.Fatalf("ensurePicoChannel() error = %v", err)
|
||||
if _, err := h.EnsurePicoChannel(""); err != nil {
|
||||
t.Fatalf("EnsurePicoChannel() error = %v", err)
|
||||
}
|
||||
|
||||
cfg, err := config.LoadConfig(configPath)
|
||||
|
|
@ -60,8 +61,8 @@ func TestEnsurePicoChannel_DoesNotSetWildcardOrigins(t *testing.T) {
|
|||
configPath := filepath.Join(t.TempDir(), "config.json")
|
||||
h := NewHandler(configPath)
|
||||
|
||||
if _, err := h.ensurePicoChannel("http://localhost:18800"); err != nil {
|
||||
t.Fatalf("ensurePicoChannel() error = %v", err)
|
||||
if _, err := h.EnsurePicoChannel("http://localhost:18800"); err != nil {
|
||||
t.Fatalf("EnsurePicoChannel() error = %v", err)
|
||||
}
|
||||
|
||||
cfg, err := config.LoadConfig(configPath)
|
||||
|
|
@ -80,8 +81,8 @@ func TestEnsurePicoChannel_NoOriginWithoutCaller(t *testing.T) {
|
|||
configPath := filepath.Join(t.TempDir(), "config.json")
|
||||
h := NewHandler(configPath)
|
||||
|
||||
if _, err := h.ensurePicoChannel(""); err != nil {
|
||||
t.Fatalf("ensurePicoChannel() error = %v", err)
|
||||
if _, err := h.EnsurePicoChannel(""); err != nil {
|
||||
t.Fatalf("EnsurePicoChannel() error = %v", err)
|
||||
}
|
||||
|
||||
cfg, err := config.LoadConfig(configPath)
|
||||
|
|
@ -101,8 +102,8 @@ func TestEnsurePicoChannel_SetsCallerOrigin(t *testing.T) {
|
|||
h := NewHandler(configPath)
|
||||
|
||||
lanOrigin := "http://192.168.1.9:18800"
|
||||
if _, err := h.ensurePicoChannel(lanOrigin); err != nil {
|
||||
t.Fatalf("ensurePicoChannel() error = %v", err)
|
||||
if _, err := h.EnsurePicoChannel(lanOrigin); err != nil {
|
||||
t.Fatalf("EnsurePicoChannel() error = %v", err)
|
||||
}
|
||||
|
||||
cfg, err := config.LoadConfig(configPath)
|
||||
|
|
@ -130,12 +131,12 @@ func TestEnsurePicoChannel_PreservesUserSettings(t *testing.T) {
|
|||
|
||||
h := NewHandler(configPath)
|
||||
|
||||
changed, err := h.ensurePicoChannel("")
|
||||
changed, err := h.EnsurePicoChannel("")
|
||||
if err != nil {
|
||||
t.Fatalf("ensurePicoChannel() error = %v", err)
|
||||
t.Fatalf("EnsurePicoChannel() error = %v", err)
|
||||
}
|
||||
if changed {
|
||||
t.Error("ensurePicoChannel() should not change a fully configured config")
|
||||
t.Error("EnsurePicoChannel() should not change a fully configured config")
|
||||
}
|
||||
|
||||
cfg, err = config.LoadConfig(configPath)
|
||||
|
|
@ -154,6 +155,71 @@ func TestEnsurePicoChannel_PreservesUserSettings(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestEnsurePicoChannel_ExistingConfigWithoutSecurityFile(t *testing.T) {
|
||||
configPath := filepath.Join(t.TempDir(), "config.json")
|
||||
|
||||
cfg := config.DefaultConfig()
|
||||
raw, err := json.Marshal(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("Marshal() error = %v", err)
|
||||
}
|
||||
if err = os.WriteFile(configPath, raw, 0o600); err != nil {
|
||||
t.Fatalf("WriteFile() error = %v", err)
|
||||
}
|
||||
|
||||
h := NewHandler(configPath)
|
||||
|
||||
changed, err := h.EnsurePicoChannel("")
|
||||
if err != nil {
|
||||
t.Fatalf("EnsurePicoChannel() error = %v", err)
|
||||
}
|
||||
if !changed {
|
||||
t.Fatal("EnsurePicoChannel() should report changed when pico is missing")
|
||||
}
|
||||
|
||||
cfg, err = config.LoadConfig(configPath)
|
||||
if err != nil {
|
||||
t.Fatalf("LoadConfig() error = %v", err)
|
||||
}
|
||||
|
||||
if !cfg.Channels.Pico.Enabled {
|
||||
t.Error("expected Pico to be enabled after setup")
|
||||
}
|
||||
if cfg.Channels.Pico.Token() == "" {
|
||||
t.Error("expected a non-empty token after setup")
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(filepath.Dir(configPath), config.SecurityConfigFile)); err != nil {
|
||||
t.Fatalf("expected .security.yml to be created: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnsurePicoChannel_ConfiguresPicoWithoutGateway(t *testing.T) {
|
||||
configPath := filepath.Join(t.TempDir(), "config.json")
|
||||
|
||||
cfg := config.DefaultConfig()
|
||||
cfg.Agents.Defaults.ModelName = ""
|
||||
if err := config.SaveConfig(configPath, cfg); err != nil {
|
||||
t.Fatalf("SaveConfig() error = %v", err)
|
||||
}
|
||||
|
||||
h := NewHandler(configPath)
|
||||
if _, err := h.EnsurePicoChannel(""); err != nil {
|
||||
t.Fatalf("EnsurePicoChannel() error = %v", err)
|
||||
}
|
||||
|
||||
cfg, err := config.LoadConfig(configPath)
|
||||
if err != nil {
|
||||
t.Fatalf("LoadConfig() error = %v", err)
|
||||
}
|
||||
|
||||
if !cfg.Channels.Pico.Enabled {
|
||||
t.Error("expected Pico to be enabled after launcher startup setup")
|
||||
}
|
||||
if cfg.Channels.Pico.Token() == "" {
|
||||
t.Error("expected a non-empty token after launcher startup setup")
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnsurePicoChannel_Idempotent(t *testing.T) {
|
||||
configPath := filepath.Join(t.TempDir(), "config.json")
|
||||
h := NewHandler(configPath)
|
||||
|
|
@ -161,20 +227,20 @@ func TestEnsurePicoChannel_Idempotent(t *testing.T) {
|
|||
origin := "http://localhost:18800"
|
||||
|
||||
// First call sets things up
|
||||
if _, err := h.ensurePicoChannel(origin); err != nil {
|
||||
t.Fatalf("first ensurePicoChannel() error = %v", err)
|
||||
if _, err := h.EnsurePicoChannel(origin); err != nil {
|
||||
t.Fatalf("first EnsurePicoChannel() error = %v", err)
|
||||
}
|
||||
|
||||
cfg1, _ := config.LoadConfig(configPath)
|
||||
token1 := cfg1.Channels.Pico.Token()
|
||||
|
||||
// Second call should be a no-op
|
||||
changed, err := h.ensurePicoChannel(origin)
|
||||
changed, err := h.EnsurePicoChannel(origin)
|
||||
if err != nil {
|
||||
t.Fatalf("second ensurePicoChannel() error = %v", err)
|
||||
t.Fatalf("second EnsurePicoChannel() error = %v", err)
|
||||
}
|
||||
if changed {
|
||||
t.Error("second ensurePicoChannel() should not report changed")
|
||||
t.Error("second EnsurePicoChannel() should not report changed")
|
||||
}
|
||||
|
||||
cfg2, _ := config.LoadConfig(configPath)
|
||||
|
|
|
|||
|
|
@ -17,15 +17,18 @@ type Handler struct {
|
|||
oauthMu sync.Mutex
|
||||
oauthFlows map[string]*oauthFlow
|
||||
oauthState map[string]string
|
||||
weixinMu sync.Mutex
|
||||
weixinFlows map[string]*weixinFlow
|
||||
}
|
||||
|
||||
// NewHandler creates an instance of the API handler.
|
||||
func NewHandler(configPath string) *Handler {
|
||||
return &Handler{
|
||||
configPath: configPath,
|
||||
serverPort: launcherconfig.DefaultPort,
|
||||
oauthFlows: make(map[string]*oauthFlow),
|
||||
oauthState: make(map[string]string),
|
||||
configPath: configPath,
|
||||
serverPort: launcherconfig.DefaultPort,
|
||||
oauthFlows: make(map[string]*oauthFlow),
|
||||
oauthState: make(map[string]string),
|
||||
weixinFlows: make(map[string]*weixinFlow),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -69,6 +72,9 @@ func (h *Handler) RegisterRoutes(mux *http.ServeMux) {
|
|||
|
||||
// Launcher service parameters (port/public)
|
||||
h.registerLauncherConfigRoutes(mux)
|
||||
|
||||
// WeChat QR login flow
|
||||
h.registerWeixinRoutes(mux)
|
||||
}
|
||||
|
||||
// Shutdown gracefully shuts down the handler, stopping the gateway if it was started by this handler.
|
||||
|
|
|
|||
317
web/backend/api/weixin.go
Normal file
317
web/backend/api/weixin.go
Normal file
|
|
@ -0,0 +1,317 @@
|
|||
package api
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/base64"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"rsc.io/qr"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/channels/weixin"
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
"github.com/sipeed/picoclaw/pkg/logger"
|
||||
)
|
||||
|
||||
const (
|
||||
weixinFlowTTL = 5 * time.Minute
|
||||
weixinFlowGCAge = 30 * time.Minute
|
||||
weixinBaseURL = "https://ilinkai.weixin.qq.com/"
|
||||
weixinBotType = "3"
|
||||
)
|
||||
|
||||
const (
|
||||
weixinStatusWait = "wait"
|
||||
weixinStatusScanned = "scaned"
|
||||
weixinStatusConfirmed = "confirmed"
|
||||
weixinStatusExpired = "expired"
|
||||
weixinStatusError = "error"
|
||||
)
|
||||
|
||||
type weixinFlow struct {
|
||||
ID string
|
||||
Qrcode string // qrcode token from WeChat API (used for status polling)
|
||||
QRDataURI string // base64 PNG data URI for display
|
||||
AccountID string // IlinkBotID returned on confirmed
|
||||
Status string // wait / scaned / confirmed / expired / error
|
||||
Error string
|
||||
CreatedAt time.Time
|
||||
UpdatedAt time.Time
|
||||
ExpiresAt time.Time
|
||||
}
|
||||
|
||||
type weixinFlowResponse struct {
|
||||
FlowID string `json:"flow_id"`
|
||||
Status string `json:"status"`
|
||||
QRDataURI string `json:"qr_data_uri,omitempty"`
|
||||
AccountID string `json:"account_id,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
// registerWeixinRoutes binds WeChat QR login endpoints to the ServeMux.
|
||||
func (h *Handler) registerWeixinRoutes(mux *http.ServeMux) {
|
||||
mux.HandleFunc("POST /api/weixin/flows", h.handleStartWeixinFlow)
|
||||
mux.HandleFunc("GET /api/weixin/flows/{id}", h.handlePollWeixinFlow)
|
||||
}
|
||||
|
||||
// handleStartWeixinFlow starts a new WeChat QR login flow.
|
||||
//
|
||||
// POST /api/weixin/flows
|
||||
func (h *Handler) handleStartWeixinFlow(w http.ResponseWriter, r *http.Request) {
|
||||
ctx, cancel := context.WithTimeout(r.Context(), 15*time.Second)
|
||||
defer cancel()
|
||||
|
||||
api, err := weixin.NewApiClient(weixinBaseURL, "", "")
|
||||
if err != nil {
|
||||
http.Error(w, fmt.Sprintf("failed to create weixin client: %v", err), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
qrResp, err := api.GetQRCode(ctx, weixinBotType)
|
||||
if err != nil {
|
||||
http.Error(w, fmt.Sprintf("failed to get QR code: %v", err), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
dataURI, err := generateQRDataURI(qrResp.QrcodeImgContent)
|
||||
if err != nil {
|
||||
http.Error(w, fmt.Sprintf("failed to generate QR image: %v", err), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
flow := &weixinFlow{
|
||||
ID: newWeixinFlowID(),
|
||||
Qrcode: qrResp.Qrcode,
|
||||
QRDataURI: dataURI,
|
||||
Status: weixinStatusWait,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
ExpiresAt: now.Add(weixinFlowTTL),
|
||||
}
|
||||
h.storeWeixinFlow(flow)
|
||||
|
||||
logger.InfoCF("weixin", "QR flow started", map[string]any{"flow_id": flow.ID})
|
||||
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(weixinFlowResponse{
|
||||
FlowID: flow.ID,
|
||||
Status: flow.Status,
|
||||
QRDataURI: flow.QRDataURI,
|
||||
})
|
||||
}
|
||||
|
||||
// handlePollWeixinFlow polls the WeChat API for QR code status and updates the flow.
|
||||
//
|
||||
// GET /api/weixin/flows/{id}
|
||||
func (h *Handler) handlePollWeixinFlow(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.getWeixinFlow(flowID)
|
||||
if !ok {
|
||||
http.Error(w, "flow not found", http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
|
||||
// Return terminal states directly without polling WeChat again
|
||||
if flow.Status == weixinStatusConfirmed ||
|
||||
flow.Status == weixinStatusExpired ||
|
||||
flow.Status == weixinStatusError {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(weixinFlowResponse{
|
||||
FlowID: flow.ID,
|
||||
Status: flow.Status,
|
||||
Error: flow.Error,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(r.Context(), 10*time.Second)
|
||||
defer cancel()
|
||||
|
||||
api, err := weixin.NewApiClient(weixinBaseURL, "", "")
|
||||
if err != nil {
|
||||
h.setWeixinFlowError(flowID, fmt.Sprintf("client error: %v", err))
|
||||
flow, _ = h.getWeixinFlow(flowID)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(weixinFlowResponse{FlowID: flow.ID, Status: flow.Status, Error: flow.Error})
|
||||
return
|
||||
}
|
||||
|
||||
statusResp, err := api.GetQRCodeStatus(ctx, flow.Qrcode)
|
||||
if err != nil {
|
||||
// Transient error — keep current status, return it
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(weixinFlowResponse{
|
||||
FlowID: flow.ID,
|
||||
Status: flow.Status,
|
||||
QRDataURI: flow.QRDataURI,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
switch statusResp.Status {
|
||||
case weixinStatusWait:
|
||||
// no change
|
||||
|
||||
case weixinStatusScanned:
|
||||
h.updateWeixinFlowStatus(flowID, weixinStatusScanned)
|
||||
|
||||
case weixinStatusConfirmed:
|
||||
if statusResp.BotToken == "" {
|
||||
h.setWeixinFlowError(flowID, "login confirmed but missing bot_token")
|
||||
break
|
||||
}
|
||||
if saveErr := h.saveWeixinBinding(statusResp.BotToken, statusResp.IlinkBotID); saveErr != nil {
|
||||
h.setWeixinFlowError(flowID, fmt.Sprintf("failed to save token: %v", saveErr))
|
||||
logger.ErrorCF("weixin", "failed to save token", map[string]any{"error": saveErr.Error()})
|
||||
break
|
||||
}
|
||||
h.setWeixinFlowConfirmed(flowID, statusResp.IlinkBotID)
|
||||
logger.InfoCF("weixin", "QR login confirmed, token saved", map[string]any{
|
||||
"flow_id": flowID,
|
||||
"account_id": statusResp.IlinkBotID,
|
||||
})
|
||||
|
||||
case weixinStatusExpired:
|
||||
h.updateWeixinFlowStatus(flowID, weixinStatusExpired)
|
||||
|
||||
default:
|
||||
// unknown status, keep as-is
|
||||
}
|
||||
|
||||
flow, _ = h.getWeixinFlow(flowID)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
resp := weixinFlowResponse{
|
||||
FlowID: flow.ID,
|
||||
Status: flow.Status,
|
||||
AccountID: flow.AccountID,
|
||||
Error: flow.Error,
|
||||
}
|
||||
if flow.Status == weixinStatusWait || flow.Status == weixinStatusScanned {
|
||||
resp.QRDataURI = flow.QRDataURI
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(resp)
|
||||
}
|
||||
|
||||
// saveWeixinBinding writes the token/account ID, enables the Weixin channel,
|
||||
// and best-effort restarts the gateway when it is currently running.
|
||||
func (h *Handler) saveWeixinBinding(token, accountID string) error {
|
||||
cfg, err := config.LoadConfig(h.configPath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("load config: %w", err)
|
||||
}
|
||||
cfg.Channels.Weixin.SetToken(token)
|
||||
cfg.Channels.Weixin.Enabled = true
|
||||
if accountID != "" {
|
||||
cfg.Channels.Weixin.AccountID = accountID
|
||||
}
|
||||
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("weixin", "failed to restart gateway after saving binding", map[string]any{
|
||||
"error": err.Error(),
|
||||
})
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// generateQRDataURI encodes content as a QR code PNG and returns a data URI.
|
||||
func generateQRDataURI(content string) (string, error) {
|
||||
code, err := qr.Encode(content, qr.L)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("qr encode: %w", err)
|
||||
}
|
||||
pngBytes := code.PNG()
|
||||
encoded := base64.StdEncoding.EncodeToString(pngBytes)
|
||||
return "data:image/png;base64," + encoded, nil
|
||||
}
|
||||
|
||||
func newWeixinFlowID() string {
|
||||
buf := make([]byte, 12)
|
||||
if _, err := rand.Read(buf); err != nil {
|
||||
return fmt.Sprintf("wx_%d", time.Now().UnixNano())
|
||||
}
|
||||
return "wx_" + hex.EncodeToString(buf)
|
||||
}
|
||||
|
||||
func (h *Handler) storeWeixinFlow(flow *weixinFlow) {
|
||||
h.weixinMu.Lock()
|
||||
defer h.weixinMu.Unlock()
|
||||
h.gcWeixinFlowsLocked(time.Now())
|
||||
h.weixinFlows[flow.ID] = flow
|
||||
}
|
||||
|
||||
func (h *Handler) getWeixinFlow(flowID string) (*weixinFlow, bool) {
|
||||
h.weixinMu.Lock()
|
||||
defer h.weixinMu.Unlock()
|
||||
h.gcWeixinFlowsLocked(time.Now())
|
||||
flow, ok := h.weixinFlows[flowID]
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
cp := *flow
|
||||
return &cp, true
|
||||
}
|
||||
|
||||
func (h *Handler) updateWeixinFlowStatus(flowID, status string) {
|
||||
h.weixinMu.Lock()
|
||||
defer h.weixinMu.Unlock()
|
||||
if flow, ok := h.weixinFlows[flowID]; ok {
|
||||
flow.Status = status
|
||||
flow.UpdatedAt = time.Now()
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Handler) setWeixinFlowConfirmed(flowID, accountID string) {
|
||||
h.weixinMu.Lock()
|
||||
defer h.weixinMu.Unlock()
|
||||
if flow, ok := h.weixinFlows[flowID]; ok {
|
||||
flow.Status = weixinStatusConfirmed
|
||||
flow.AccountID = accountID
|
||||
flow.UpdatedAt = time.Now()
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Handler) setWeixinFlowError(flowID, errMsg string) {
|
||||
h.weixinMu.Lock()
|
||||
defer h.weixinMu.Unlock()
|
||||
if flow, ok := h.weixinFlows[flowID]; ok {
|
||||
flow.Status = weixinStatusError
|
||||
flow.Error = errMsg
|
||||
flow.UpdatedAt = time.Now()
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Handler) gcWeixinFlowsLocked(now time.Time) {
|
||||
for id, flow := range h.weixinFlows {
|
||||
if flow.Status == weixinStatusWait || flow.Status == weixinStatusScanned {
|
||||
if !flow.ExpiresAt.IsZero() && now.After(flow.ExpiresAt) {
|
||||
flow.Status = weixinStatusExpired
|
||||
flow.UpdatedAt = now
|
||||
}
|
||||
}
|
||||
if flow.Status != weixinStatusWait &&
|
||||
flow.Status != weixinStatusScanned &&
|
||||
now.Sub(flow.UpdatedAt) > weixinFlowGCAge {
|
||||
delete(h.weixinFlows, id)
|
||||
}
|
||||
}
|
||||
}
|
||||
56
web/backend/api/weixin_test.go
Normal file
56
web/backend/api/weixin_test.go
Normal file
|
|
@ -0,0 +1,56 @@
|
|||
package api
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
)
|
||||
|
||||
func TestSaveWeixinBindingReturnsSuccessWhenRestartFails(t *testing.T) {
|
||||
resetGatewayTestState(t)
|
||||
|
||||
configPath := filepath.Join(t.TempDir(), "config.json")
|
||||
cfg := config.DefaultConfig()
|
||||
if err := config.SaveConfig(configPath, cfg); err != nil {
|
||||
t.Fatalf("SaveConfig() error = %v", err)
|
||||
}
|
||||
|
||||
originalHealthGet := gatewayHealthGet
|
||||
gatewayHealthGet = func(url string, timeout time.Duration) (*http.Response, error) {
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Body: io.NopCloser(strings.NewReader(
|
||||
`{"status":"ok","uptime":"1s","pid":` + strconv.Itoa(os.Getpid()) + `}`,
|
||||
)),
|
||||
}, nil
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
gatewayHealthGet = originalHealthGet
|
||||
})
|
||||
|
||||
h := NewHandler(configPath)
|
||||
if err := h.saveWeixinBinding("bot-token", "bot-account"); err != nil {
|
||||
t.Fatalf("saveWeixinBinding() error = %v, want nil after config save succeeds", err)
|
||||
}
|
||||
|
||||
savedCfg, err := config.LoadConfig(configPath)
|
||||
if err != nil {
|
||||
t.Fatalf("LoadConfig() error = %v", err)
|
||||
}
|
||||
if got := savedCfg.Channels.Weixin.Token(); got != "bot-token" {
|
||||
t.Fatalf("Weixin.Token() = %q, want %q", got, "bot-token")
|
||||
}
|
||||
if got := savedCfg.Channels.Weixin.AccountID; got != "bot-account" {
|
||||
t.Fatalf("Weixin.AccountID = %q, want %q", got, "bot-account")
|
||||
}
|
||||
if !savedCfg.Channels.Weixin.Enabled {
|
||||
t.Fatalf("Weixin.Enabled = false, want true")
|
||||
}
|
||||
}
|
||||
|
|
@ -169,6 +169,9 @@ func main() {
|
|||
|
||||
// API Routes (e.g. /api/status)
|
||||
apiHandler = api.NewHandler(absPath)
|
||||
if _, err = apiHandler.EnsurePicoChannel(""); err != nil {
|
||||
logger.ErrorC("web", fmt.Sprintf("Warning: failed to ensure pico channel on startup: %v", err))
|
||||
}
|
||||
apiHandler.SetServerOptions(portNum, effectivePublic, explicitPublic, launcherCfg.AllowedCIDRs)
|
||||
apiHandler.RegisterRoutes(mux)
|
||||
|
||||
|
|
|
|||
|
|
@ -62,4 +62,26 @@ export async function patchAppConfig(
|
|||
})
|
||||
}
|
||||
|
||||
// WeChat QR login flow API
|
||||
|
||||
export interface WeixinFlowResponse {
|
||||
flow_id: string
|
||||
status: "wait" | "scaned" | "confirmed" | "expired" | "error"
|
||||
qr_data_uri?: string
|
||||
account_id?: string
|
||||
error?: string
|
||||
}
|
||||
|
||||
export async function startWeixinFlow(): Promise<WeixinFlowResponse> {
|
||||
return request<WeixinFlowResponse>("/api/weixin/flows", { method: "POST" })
|
||||
}
|
||||
|
||||
export async function pollWeixinFlow(
|
||||
flowID: string,
|
||||
): Promise<WeixinFlowResponse> {
|
||||
return request<WeixinFlowResponse>(
|
||||
`/api/weixin/flows/${encodeURIComponent(flowID)}`,
|
||||
)
|
||||
}
|
||||
|
||||
export type { ChannelsCatalogResponse, ConfigActionResponse }
|
||||
|
|
|
|||
|
|
@ -67,14 +67,17 @@ const baseNavGroups: Omit<NavGroup, "items">[] = [
|
|||
|
||||
export function AppSidebar({ ...props }: React.ComponentProps<typeof Sidebar>) {
|
||||
const routerState = useRouterState()
|
||||
const { t } = useTranslation()
|
||||
const { i18n, t } = useTranslation()
|
||||
const currentPath = routerState.location.pathname
|
||||
const {
|
||||
channelItems,
|
||||
hasMoreChannels,
|
||||
showAllChannels,
|
||||
toggleShowAllChannels,
|
||||
} = useSidebarChannels({ t })
|
||||
} = useSidebarChannels({
|
||||
language: (i18n.resolvedLanguage ?? i18n.language ?? "").toLowerCase(),
|
||||
t,
|
||||
})
|
||||
|
||||
const navGroups: NavGroup[] = React.useMemo(() => {
|
||||
return [
|
||||
|
|
|
|||
|
|
@ -1,8 +1,6 @@
|
|||
import { IconLoader2 } from "@tabler/icons-react"
|
||||
import { useAtomValue } from "jotai"
|
||||
import { useCallback, useEffect, useMemo, useRef, useState } from "react"
|
||||
import { useTranslation } from "react-i18next"
|
||||
import { toast } from "sonner"
|
||||
|
||||
import {
|
||||
type ChannelConfig,
|
||||
|
|
@ -17,10 +15,12 @@ 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 { WeixinForm } from "@/components/channels/channel-forms/weixin-form"
|
||||
import { PageHeader } from "@/components/page-header"
|
||||
import { Button } from "@/components/ui/button"
|
||||
import { Switch } from "@/components/ui/switch"
|
||||
import { gatewayAtom } from "@/store/gateway"
|
||||
import { useGateway } from "@/hooks/use-gateway"
|
||||
import { refreshGatewayState } from "@/store/gateway"
|
||||
|
||||
interface ChannelConfigPageProps {
|
||||
channelName: string
|
||||
|
|
@ -142,14 +142,10 @@ function isConfigured(
|
|||
)
|
||||
case "onebot":
|
||||
return asString(config.ws_url) !== ""
|
||||
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":
|
||||
|
|
@ -190,11 +186,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 ["bot_id", "secret"]
|
||||
case "whatsapp":
|
||||
return ["bridge_url"]
|
||||
case "pico":
|
||||
|
|
@ -238,7 +230,7 @@ const CHANNELS_WITHOUT_DOCS = new Set([
|
|||
|
||||
export function ChannelConfigPage({ channelName }: ChannelConfigPageProps) {
|
||||
const { t, i18n } = useTranslation()
|
||||
const gateway = useAtomValue(gatewayAtom)
|
||||
const { state: gatewayState } = useGateway()
|
||||
|
||||
const [loading, setLoading] = useState(true)
|
||||
const [saving, setSaving] = useState(false)
|
||||
|
|
@ -251,56 +243,59 @@ export function ChannelConfigPage({ channelName }: ChannelConfigPageProps) {
|
|||
const [editConfig, setEditConfig] = useState<ChannelConfig>({})
|
||||
const [enabled, setEnabled] = useState(false)
|
||||
|
||||
const loadData = useCallback(async () => {
|
||||
setLoading(true)
|
||||
try {
|
||||
const [catalog, appConfig] = await Promise.all([
|
||||
getChannelsCatalog(),
|
||||
getAppConfig(),
|
||||
])
|
||||
const matched =
|
||||
catalog.channels.find((item) => item.name === channelName) ?? null
|
||||
const loadData = useCallback(
|
||||
async (silent = false) => {
|
||||
if (!silent) setLoading(true)
|
||||
try {
|
||||
const [catalog, appConfig] = await Promise.all([
|
||||
getChannelsCatalog(),
|
||||
getAppConfig(),
|
||||
])
|
||||
const matched =
|
||||
catalog.channels.find((item) => item.name === channelName) ?? null
|
||||
|
||||
if (!matched) {
|
||||
setChannel(null)
|
||||
setFetchError(
|
||||
t("channels.page.notFound", {
|
||||
name: channelName,
|
||||
}),
|
||||
)
|
||||
return
|
||||
if (!matched) {
|
||||
setChannel(null)
|
||||
setFetchError(
|
||||
t("channels.page.notFound", {
|
||||
name: channelName,
|
||||
}),
|
||||
)
|
||||
return
|
||||
}
|
||||
|
||||
const channelsConfig = asRecord(asRecord(appConfig).channels)
|
||||
const raw = asRecord(channelsConfig[matched.config_key])
|
||||
const normalized = normalizeConfig(matched, raw)
|
||||
|
||||
setChannel(matched)
|
||||
setBaseConfig(normalized)
|
||||
setEditConfig(buildEditConfig(normalized))
|
||||
setEnabled(asBool(normalized.enabled))
|
||||
setFetchError("")
|
||||
setServerError("")
|
||||
setFieldErrors({})
|
||||
} catch (e) {
|
||||
setFetchError(e instanceof Error ? e.message : t("channels.loadError"))
|
||||
} finally {
|
||||
if (!silent) setLoading(false)
|
||||
}
|
||||
|
||||
const channelsConfig = asRecord(asRecord(appConfig).channels)
|
||||
const raw = asRecord(channelsConfig[matched.config_key])
|
||||
const normalized = normalizeConfig(matched, raw)
|
||||
|
||||
setChannel(matched)
|
||||
setBaseConfig(normalized)
|
||||
setEditConfig(buildEditConfig(normalized))
|
||||
setEnabled(asBool(normalized.enabled))
|
||||
setFetchError("")
|
||||
setServerError("")
|
||||
setFieldErrors({})
|
||||
} catch (e) {
|
||||
setFetchError(e instanceof Error ? e.message : t("channels.loadError"))
|
||||
} finally {
|
||||
setLoading(false)
|
||||
}
|
||||
}, [channelName, t])
|
||||
},
|
||||
[channelName, t],
|
||||
)
|
||||
|
||||
useEffect(() => {
|
||||
loadData()
|
||||
}, [loadData])
|
||||
|
||||
const previousGatewayStatusRef = useRef(gateway.status)
|
||||
const previousGatewayStatusRef = useRef(gatewayState)
|
||||
useEffect(() => {
|
||||
const previousStatus = previousGatewayStatusRef.current
|
||||
if (previousStatus !== "running" && gateway.status === "running") {
|
||||
if (previousStatus !== "running" && gatewayState === "running") {
|
||||
void loadData()
|
||||
}
|
||||
previousGatewayStatusRef.current = gateway.status
|
||||
}, [gateway.status, loadData])
|
||||
previousGatewayStatusRef.current = gatewayState
|
||||
}, [gatewayState, loadData])
|
||||
|
||||
const savePayload = useMemo(() => {
|
||||
if (!channel) return null
|
||||
|
|
@ -393,18 +388,28 @@ export function ChannelConfigPage({ channelName }: ChannelConfigPageProps) {
|
|||
[channel.config_key]: savePayload,
|
||||
},
|
||||
})
|
||||
toast.success(t("channels.page.saveSuccess"))
|
||||
await loadData()
|
||||
} catch (e) {
|
||||
const message =
|
||||
e instanceof Error ? e.message : t("channels.page.saveError")
|
||||
setServerError(message)
|
||||
toast.error(message)
|
||||
} finally {
|
||||
setSaving(false)
|
||||
}
|
||||
}
|
||||
|
||||
const handleWeixinBindSuccess = 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 renderForm = () => {
|
||||
if (!channel) return null
|
||||
const isEdit = configured
|
||||
|
|
@ -446,6 +451,15 @@ export function ChannelConfigPage({ channelName }: ChannelConfigPageProps) {
|
|||
fieldErrors={fieldErrors}
|
||||
/>
|
||||
)
|
||||
case "weixin":
|
||||
return (
|
||||
<WeixinForm
|
||||
config={editConfig}
|
||||
onChange={handleChange}
|
||||
isEdit={isEdit}
|
||||
onBindSuccess={() => void handleWeixinBindSuccess()}
|
||||
/>
|
||||
)
|
||||
default:
|
||||
return (
|
||||
<GenericForm
|
||||
|
|
|
|||
|
|
@ -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,14 @@ 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"),
|
||||
|
|
|
|||
|
|
@ -0,0 +1,351 @@
|
|||
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 { pollWeixinFlow, startWeixinFlow } from "@/api/channels"
|
||||
import { Field } from "@/components/shared-form"
|
||||
import { Button } from "@/components/ui/button"
|
||||
import { Input } from "@/components/ui/input"
|
||||
|
||||
type BindingState =
|
||||
| "idle"
|
||||
| "loading"
|
||||
| "waiting"
|
||||
| "scaned"
|
||||
| "confirmed"
|
||||
| "expired"
|
||||
| "error"
|
||||
|
||||
interface WeixinFormProps {
|
||||
config: ChannelConfig
|
||||
onChange: (key: string, value: unknown) => void
|
||||
isEdit: boolean
|
||||
onBindSuccess?: () => void
|
||||
}
|
||||
|
||||
function asString(value: unknown): string {
|
||||
return typeof value === "string" ? value : ""
|
||||
}
|
||||
|
||||
function asStringArray(value: unknown): string[] {
|
||||
if (!Array.isArray(value)) return []
|
||||
return value.filter((item): item is string => typeof item === "string")
|
||||
}
|
||||
|
||||
export function WeixinForm({
|
||||
config,
|
||||
onChange,
|
||||
isEdit,
|
||||
onBindSuccess,
|
||||
}: WeixinFormProps) {
|
||||
const { t } = useTranslation()
|
||||
|
||||
const [bindState, setBindState] = useState<BindingState>("idle")
|
||||
const [qrDataURI, setQrDataURI] = useState<string | null>(null)
|
||||
const [accountID, setAccountID] = useState<string | null>(null)
|
||||
const [errorMsg, setErrorMsg] = useState("")
|
||||
|
||||
const pollTimerRef = useRef<ReturnType<typeof setInterval> | null>(null)
|
||||
const pollGenerationRef = useRef(0)
|
||||
const isBound = isEdit && asString(config.account_id) !== ""
|
||||
const existingAccountID = asString(config.account_id)
|
||||
|
||||
const stopPolling = useCallback(() => {
|
||||
pollGenerationRef.current += 1
|
||||
if (pollTimerRef.current !== null) {
|
||||
clearInterval(pollTimerRef.current)
|
||||
pollTimerRef.current = null
|
||||
}
|
||||
}, [])
|
||||
|
||||
useEffect(() => () => stopPolling(), [stopPolling])
|
||||
|
||||
useEffect(() => {
|
||||
if (!existingAccountID) return
|
||||
stopPolling()
|
||||
setAccountID(existingAccountID)
|
||||
setBindState("confirmed")
|
||||
setErrorMsg("")
|
||||
}, [existingAccountID, 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 pollWeixinFlow(id)
|
||||
if (generation !== pollGenerationRef.current) {
|
||||
return
|
||||
}
|
||||
if (resp.status === "scaned") {
|
||||
setBindState("scaned")
|
||||
} else if (resp.status === "confirmed") {
|
||||
stopPolling()
|
||||
setAccountID(resp.account_id ?? existingAccountID ?? 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.weixin.errorGeneric"))
|
||||
}
|
||||
} catch {
|
||||
// transient network error — keep polling
|
||||
} finally {
|
||||
inFlight = false
|
||||
}
|
||||
}, 2000)
|
||||
},
|
||||
[existingAccountID, stopPolling, onBindSuccess, t],
|
||||
)
|
||||
|
||||
const handleBind = async () => {
|
||||
setBindState("loading")
|
||||
setErrorMsg("")
|
||||
setQrDataURI(null)
|
||||
stopPolling()
|
||||
try {
|
||||
const resp = await startWeixinFlow()
|
||||
setQrDataURI(resp.qr_data_uri ?? null)
|
||||
setBindState("waiting")
|
||||
startPolling(resp.flow_id)
|
||||
} catch (e) {
|
||||
setBindState("error")
|
||||
setErrorMsg(
|
||||
e instanceof Error ? e.message : t("channels.weixin.errorGeneric"),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
const handleRebind = () => {
|
||||
stopPolling()
|
||||
setBindState("idle")
|
||||
setQrDataURI(null)
|
||||
setAccountID(null)
|
||||
setErrorMsg("")
|
||||
void handleBind()
|
||||
}
|
||||
|
||||
const renderBindSection = () => {
|
||||
if (bindState === "idle") {
|
||||
if (isBound) {
|
||||
return (
|
||||
<div className="flex flex-col items-center gap-3 py-6">
|
||||
<div className="flex items-center gap-2 rounded-full bg-emerald-500/10 px-4 py-2 text-sm font-medium text-emerald-600 dark:text-emerald-400">
|
||||
<IconCheck size={16} />
|
||||
{t("channels.weixin.bound")}
|
||||
</div>
|
||||
{existingAccountID && (
|
||||
<p className="text-muted-foreground font-mono text-xs">
|
||||
{existingAccountID}
|
||||
</p>
|
||||
)}
|
||||
<Button
|
||||
variant="outline"
|
||||
size="sm"
|
||||
onClick={handleRebind}
|
||||
className="mt-1 gap-2"
|
||||
>
|
||||
<IconRefresh size={14} />
|
||||
{t("channels.weixin.rebind")}
|
||||
</Button>
|
||||
</div>
|
||||
)
|
||||
}
|
||||
return (
|
||||
<div className="flex flex-col items-center gap-4 py-6">
|
||||
<p className="text-muted-foreground text-sm">
|
||||
{t("channels.weixin.notBound")}
|
||||
</p>
|
||||
<Button onClick={handleBind} className="gap-2">
|
||||
<IconQrcode size={16} />
|
||||
{t("channels.weixin.bind")}
|
||||
</Button>
|
||||
</div>
|
||||
)
|
||||
}
|
||||
|
||||
if (bindState === "loading") {
|
||||
return (
|
||||
<div className="flex flex-col items-center gap-3 py-8">
|
||||
<IconLoader2
|
||||
className="text-muted-foreground animate-spin"
|
||||
size={32}
|
||||
/>
|
||||
<p className="text-muted-foreground text-sm">
|
||||
{t("channels.weixin.generating")}
|
||||
</p>
|
||||
</div>
|
||||
)
|
||||
}
|
||||
|
||||
if (bindState === "waiting" || bindState === "scaned") {
|
||||
return (
|
||||
<div className="flex flex-col items-center gap-4 py-4">
|
||||
{qrDataURI ? (
|
||||
<img
|
||||
src={qrDataURI}
|
||||
alt="WeChat QR Code"
|
||||
className="border-border/60 h-48 w-48 rounded-xl border bg-white p-2 shadow-sm"
|
||||
/>
|
||||
) : (
|
||||
<div className="border-border/60 bg-muted flex h-48 w-48 items-center justify-center rounded-xl border">
|
||||
<IconLoader2
|
||||
className="text-muted-foreground animate-spin"
|
||||
size={32}
|
||||
/>
|
||||
</div>
|
||||
)}
|
||||
{bindState === "scaned" ? (
|
||||
<div className="flex items-center gap-2 rounded-full bg-amber-500/10 px-4 py-2 text-sm font-medium text-amber-600 dark:text-amber-400">
|
||||
<IconLoader2 size={14} className="animate-spin" />
|
||||
{t("channels.weixin.scanned")}
|
||||
</div>
|
||||
) : (
|
||||
<p className="text-muted-foreground text-sm">
|
||||
{t("channels.weixin.scanHint")}
|
||||
</p>
|
||||
)}
|
||||
<Button
|
||||
variant="ghost"
|
||||
size="sm"
|
||||
onClick={handleRebind}
|
||||
className="text-muted-foreground"
|
||||
>
|
||||
<IconRefresh size={14} className="mr-1" />
|
||||
{t("channels.weixin.refresh")}
|
||||
</Button>
|
||||
</div>
|
||||
)
|
||||
}
|
||||
|
||||
if (bindState === "confirmed") {
|
||||
return (
|
||||
<div className="flex flex-col items-center gap-3 py-6">
|
||||
<div className="flex h-14 w-14 items-center justify-center rounded-full bg-emerald-500/10">
|
||||
<IconCheck
|
||||
size={28}
|
||||
className="text-emerald-600 dark:text-emerald-400"
|
||||
/>
|
||||
</div>
|
||||
<p className="text-sm font-medium text-emerald-600 dark:text-emerald-400">
|
||||
{t("channels.weixin.bound")}
|
||||
</p>
|
||||
{accountID && (
|
||||
<p className="text-muted-foreground font-mono text-xs">
|
||||
{accountID}
|
||||
</p>
|
||||
)}
|
||||
<Button
|
||||
variant="outline"
|
||||
size="sm"
|
||||
onClick={handleRebind}
|
||||
className="mt-1 gap-2"
|
||||
>
|
||||
<IconRefresh size={14} />
|
||||
{t("channels.weixin.rebind")}
|
||||
</Button>
|
||||
</div>
|
||||
)
|
||||
}
|
||||
|
||||
if (bindState === "expired") {
|
||||
return (
|
||||
<div className="flex flex-col items-center gap-4 py-6">
|
||||
<div className="flex h-14 w-14 items-center justify-center rounded-full bg-amber-500/10">
|
||||
<IconX size={28} className="text-amber-600 dark:text-amber-400" />
|
||||
</div>
|
||||
<p className="text-sm text-amber-600 dark:text-amber-400">
|
||||
{t("channels.weixin.expired")}
|
||||
</p>
|
||||
<Button onClick={handleRebind} className="gap-2">
|
||||
<IconRefresh size={14} />
|
||||
{t("channels.weixin.retry")}
|
||||
</Button>
|
||||
</div>
|
||||
)
|
||||
}
|
||||
|
||||
if (bindState === "error") {
|
||||
return (
|
||||
<div className="flex flex-col items-center gap-4 py-6">
|
||||
<div className="bg-destructive/10 flex h-14 w-14 items-center justify-center rounded-full">
|
||||
<IconX size={28} className="text-destructive" />
|
||||
</div>
|
||||
<p className="text-destructive text-sm">
|
||||
{errorMsg || t("channels.weixin.errorGeneric")}
|
||||
</p>
|
||||
<Button variant="outline" onClick={handleRebind} className="gap-2">
|
||||
<IconRefresh size={14} />
|
||||
{t("channels.weixin.retry")}
|
||||
</Button>
|
||||
</div>
|
||||
)
|
||||
}
|
||||
|
||||
return null
|
||||
}
|
||||
|
||||
return (
|
||||
<div className="space-y-5">
|
||||
{/* QR Bind Section */}
|
||||
<div className="border-border/60 bg-muted/30 rounded-xl border">
|
||||
<div className="border-border/60 border-b px-4 py-3">
|
||||
<p className="text-sm font-medium">
|
||||
{t("channels.weixin.bindTitle")}
|
||||
</p>
|
||||
<p className="text-muted-foreground mt-0.5 text-xs">
|
||||
{t("channels.weixin.bindDesc")}
|
||||
</p>
|
||||
</div>
|
||||
{renderBindSection()}
|
||||
</div>
|
||||
|
||||
{/* allow_from */}
|
||||
<Field
|
||||
label={t("channels.field.allowFrom")}
|
||||
hint={t("channels.form.desc.allowFrom")}
|
||||
>
|
||||
<Input
|
||||
value={asStringArray(config.allow_from).join(", ")}
|
||||
onChange={(e) =>
|
||||
onChange(
|
||||
"allow_from",
|
||||
e.target.value
|
||||
.split(",")
|
||||
.map((s: string) => s.trim())
|
||||
.filter(Boolean),
|
||||
)
|
||||
}
|
||||
placeholder={t("channels.field.allowFromPlaceholder")}
|
||||
/>
|
||||
</Field>
|
||||
|
||||
{/* proxy */}
|
||||
<Field
|
||||
label={t("channels.field.proxy")}
|
||||
hint={t("channels.form.desc.proxy")}
|
||||
>
|
||||
<Input
|
||||
value={asString(config.proxy)}
|
||||
onChange={(e) => onChange("proxy", e.target.value)}
|
||||
placeholder="http://localhost:7890"
|
||||
/>
|
||||
</Field>
|
||||
</div>
|
||||
)
|
||||
}
|
||||
|
|
@ -5,7 +5,7 @@ interface UserMessageProps {
|
|||
export function UserMessage({ content }: UserMessageProps) {
|
||||
return (
|
||||
<div className="flex w-full flex-col items-end gap-1.5">
|
||||
<div className="max-w-[70%] rounded-2xl rounded-tr-sm bg-violet-500 px-5 py-3 text-[15px] leading-relaxed text-white shadow-sm whitespace-pre-wrap">
|
||||
<div className="max-w-[70%] rounded-2xl rounded-tr-sm bg-violet-500 px-5 py-3 text-[15px] leading-relaxed whitespace-pre-wrap text-white shadow-sm">
|
||||
{content}
|
||||
</div>
|
||||
</div>
|
||||
|
|
|
|||
|
|
@ -179,7 +179,10 @@ export function buildFormFromConfig(config: unknown): CoreConfigForm {
|
|||
EMPTY_FORM.cronExecTimeoutMinutes,
|
||||
),
|
||||
maxTokens: asNumberString(defaults.max_tokens, EMPTY_FORM.maxTokens),
|
||||
contextWindow: asNumberString(defaults.context_window, EMPTY_FORM.contextWindow),
|
||||
contextWindow: asNumberString(
|
||||
defaults.context_window,
|
||||
EMPTY_FORM.contextWindow,
|
||||
),
|
||||
maxToolIterations: asNumberString(
|
||||
defaults.max_tool_iterations,
|
||||
EMPTY_FORM.maxToolIterations,
|
||||
|
|
|
|||
|
|
@ -28,15 +28,10 @@ import { getChannelDisplayName } from "@/components/channels/channel-display-nam
|
|||
import { gatewayAtom } from "@/store/gateway"
|
||||
|
||||
const DEFAULT_VISIBLE_CHANNELS = 4
|
||||
const CHANNEL_IMPORTANCE_ORDER = [
|
||||
"discord",
|
||||
"feishu",
|
||||
"telegram",
|
||||
const CHANNEL_IMPORTANCE_TAIL = [
|
||||
"slack",
|
||||
"line",
|
||||
"wecom",
|
||||
"wecom_app",
|
||||
"wecom_aibot",
|
||||
"dingtalk",
|
||||
"qq",
|
||||
"onebot",
|
||||
|
|
@ -47,9 +42,13 @@ const CHANNEL_IMPORTANCE_ORDER = [
|
|||
"whatsapp",
|
||||
"whatsapp_native",
|
||||
]
|
||||
const CHANNEL_IMPORTANCE_INDEX = new Map(
|
||||
CHANNEL_IMPORTANCE_ORDER.map((name, index) => [name, index]),
|
||||
)
|
||||
|
||||
function getChannelImportanceOrder(language: string): string[] {
|
||||
const priority = language.startsWith("zh")
|
||||
? ["feishu", "weixin", "discord", "telegram"]
|
||||
: ["discord", "telegram", "feishu", "weixin"]
|
||||
return [...priority, ...CHANNEL_IMPORTANCE_TAIL]
|
||||
}
|
||||
|
||||
function IconLark({ className }: { className?: string }) {
|
||||
return React.createElement("span", {
|
||||
|
|
@ -75,9 +74,8 @@ const CHANNEL_ICON_MAP: Record<
|
|||
dingtalk: IconBrandDingtalk,
|
||||
line: IconBrandLine,
|
||||
qq: IconBrandQq,
|
||||
weixin: IconBrandWechat,
|
||||
wecom: IconBrandWechat,
|
||||
wecom_app: IconBrandWechat,
|
||||
wecom_aibot: IconBrandWechat,
|
||||
whatsapp: IconBrandWhatsapp,
|
||||
whatsapp_native: IconBrandWhatsapp,
|
||||
matrix: IconBrandMatrix,
|
||||
|
|
@ -134,10 +132,11 @@ export interface SidebarChannelNavItem {
|
|||
}
|
||||
|
||||
interface UseSidebarChannelsOptions {
|
||||
language: string
|
||||
t: TFunction
|
||||
}
|
||||
|
||||
export function useSidebarChannels({ t }: UseSidebarChannelsOptions) {
|
||||
export function useSidebarChannels({ language, t }: UseSidebarChannelsOptions) {
|
||||
const gateway = useAtomValue(gatewayAtom)
|
||||
const [channels, setChannels] = React.useState<SupportedChannel[]>([])
|
||||
const [enabledMap, setEnabledMap] = React.useState<Record<string, boolean>>(
|
||||
|
|
@ -183,6 +182,12 @@ export function useSidebarChannels({ t }: UseSidebarChannelsOptions) {
|
|||
previousGatewayStatusRef.current = gateway.status
|
||||
}, [gateway.status, reloadChannels])
|
||||
|
||||
const channelImportanceIndex = React.useMemo(() => {
|
||||
return new Map(
|
||||
getChannelImportanceOrder(language).map((name, index) => [name, index]),
|
||||
)
|
||||
}, [language])
|
||||
|
||||
const sortedChannels = React.useMemo(() => {
|
||||
const list = [...channels]
|
||||
list.sort((a, b) => {
|
||||
|
|
@ -193,9 +198,9 @@ export function useSidebarChannels({ t }: UseSidebarChannelsOptions) {
|
|||
}
|
||||
|
||||
const aImportance =
|
||||
CHANNEL_IMPORTANCE_INDEX.get(a.name) ?? Number.MAX_SAFE_INTEGER
|
||||
channelImportanceIndex.get(a.name) ?? Number.MAX_SAFE_INTEGER
|
||||
const bImportance =
|
||||
CHANNEL_IMPORTANCE_INDEX.get(b.name) ?? Number.MAX_SAFE_INTEGER
|
||||
channelImportanceIndex.get(b.name) ?? Number.MAX_SAFE_INTEGER
|
||||
if (aImportance !== bImportance) {
|
||||
return aImportance - bImportance
|
||||
}
|
||||
|
|
@ -205,7 +210,7 @@ export function useSidebarChannels({ t }: UseSidebarChannelsOptions) {
|
|||
)
|
||||
})
|
||||
return list
|
||||
}, [channels, enabledMap, t])
|
||||
}, [channelImportanceIndex, channels, enabledMap, t])
|
||||
|
||||
const hasMoreChannels = sortedChannels.length > DEFAULT_VISIBLE_CHANNELS
|
||||
const visibleChannels = showAllChannels
|
||||
|
|
|
|||
|
|
@ -233,14 +233,31 @@
|
|||
"qq": "QQ",
|
||||
"onebot": "OneBot",
|
||||
"wecom": "WeCom",
|
||||
"wecom_app": "WeCom App",
|
||||
"wecom_aibot": "WeCom AI Bot",
|
||||
"whatsapp": "WhatsApp",
|
||||
"whatsapp_native": "WhatsApp Native",
|
||||
"pico": "Web",
|
||||
"maixcam": "MaixCam",
|
||||
"matrix": "Matrix",
|
||||
"irc": "IRC"
|
||||
"irc": "IRC",
|
||||
"weixin": "WeChat"
|
||||
},
|
||||
"weixin": {
|
||||
"warningTitle": "Testing phase, use with caution",
|
||||
"warningDesc": "The WeChat channel is still experimental and may carry a risk of account suspension. Use it only if you understand and accept the risk.",
|
||||
"bindEnableSuccess": "WeChat connected and the channel has been enabled automatically.",
|
||||
"bindTitle": "WeChat Account Binding",
|
||||
"bindDesc": "Scan the QR code with WeChat to bind your personal account.",
|
||||
"bind": "Bind WeChat",
|
||||
"rebind": "Re-bind",
|
||||
"bound": "WeChat Bound",
|
||||
"notBound": "WeChat account not bound yet.",
|
||||
"generating": "Generating QR code...",
|
||||
"scanHint": "Open WeChat and scan the QR code",
|
||||
"scanned": "Scanned — please confirm in WeChat",
|
||||
"expired": "QR code expired",
|
||||
"retry": "Try Again",
|
||||
"refresh": "Refresh QR",
|
||||
"errorGeneric": "An error occurred. Please try again."
|
||||
},
|
||||
"field": {
|
||||
"token": "Bot Token",
|
||||
|
|
@ -273,7 +290,9 @@
|
|||
"saveError": "Failed to save channel configuration",
|
||||
"enabled": "enabled",
|
||||
"docLink": "Documentation",
|
||||
"enableLabel": "Enable channel"
|
||||
"enableLabel": "Enable channel",
|
||||
"restartRequiredTitle": "Gateway restart required",
|
||||
"restartRequiredDesc": "The latest {{name}} configuration has been saved. Restart the gateway for it to take effect."
|
||||
},
|
||||
"form": {
|
||||
"desc": {
|
||||
|
|
|
|||
|
|
@ -233,14 +233,31 @@
|
|||
"qq": "QQ",
|
||||
"onebot": "OneBot",
|
||||
"wecom": "企业微信",
|
||||
"wecom_app": "企业微信应用",
|
||||
"wecom_aibot": "企业微信 AI 机器人",
|
||||
"whatsapp": "WhatsApp",
|
||||
"whatsapp_native": "WhatsApp Native",
|
||||
"pico": "Web",
|
||||
"maixcam": "MaixCam",
|
||||
"matrix": "Matrix",
|
||||
"irc": "IRC"
|
||||
"irc": "IRC",
|
||||
"weixin": "微信"
|
||||
},
|
||||
"weixin": {
|
||||
"warningTitle": "测试阶段,请谨慎使用",
|
||||
"warningDesc": "微信 Channel 当前仍处于测试阶段,存在封号风险。请仅在充分了解风险的前提下使用。",
|
||||
"bindEnableSuccess": "微信已连接,频道已自动启用。",
|
||||
"bindTitle": "微信账号绑定",
|
||||
"bindDesc": "使用微信扫描二维码以绑定您的个人微信账号。",
|
||||
"bind": "绑定微信",
|
||||
"rebind": "重新绑定",
|
||||
"bound": "微信已绑定",
|
||||
"notBound": "尚未绑定微信账号。",
|
||||
"generating": "正在生成二维码...",
|
||||
"scanHint": "打开微信,扫描二维码",
|
||||
"scanned": "已扫码 — 请在微信中确认",
|
||||
"expired": "二维码已过期",
|
||||
"retry": "重试",
|
||||
"refresh": "刷新二维码",
|
||||
"errorGeneric": "发生错误,请重试。"
|
||||
},
|
||||
"field": {
|
||||
"token": "Bot Token",
|
||||
|
|
@ -273,7 +290,9 @@
|
|||
"saveError": "保存频道配置失败",
|
||||
"enabled": "已启用",
|
||||
"docLink": "配置文档",
|
||||
"enableLabel": "启用频道"
|
||||
"enableLabel": "启用频道",
|
||||
"restartRequiredTitle": "需要重启服务",
|
||||
"restartRequiredDesc": "{{name}} 的最新配置已保存。重启服务后才能正式生效。"
|
||||
},
|
||||
"form": {
|
||||
"desc": {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue