azure skills whitelisting: fix skills loader, security config, and tests

This commit is contained in:
stevef 2026-03-24 11:27:04 +01:00
parent e0b6259f32
commit e93c2b2232
19 changed files with 252 additions and 177 deletions

View file

@ -273,7 +273,7 @@ test: generate
## fmt: Format Go code
fmt:
@$(GOLANGCI_LINT) fmt
@gofmt -s -w $$(find . -name "*.go" -not -path "./web/*" -not -path "./vendor/*")
## lint: Run linters
lint:

View file

@ -43,7 +43,7 @@ func NewSkillsCommand() *cobra.Command {
globalDir := filepath.Dir(internal.GetConfigPath())
globalSkillsDir := filepath.Join(globalDir, "skills")
builtinSkillsDir := filepath.Join(globalDir, "picoclaw", "skills")
d.skillsLoader = skills.NewSkillsLoader(d.workspace, globalSkillsDir, builtinSkillsDir)
d.skillsLoader = skills.NewSkillsLoader(d.workspace, globalSkillsDir, builtinSkillsDir, nil, false)
return nil
},

View file

@ -73,7 +73,7 @@ func NewContextBuilder(workspace string) *ContextBuilder {
return &ContextBuilder{
workspace: workspace,
skillsLoader: skills.NewSkillsLoader(workspace, globalSkillsDir, builtinSkillsDir),
skillsLoader: skills.NewSkillsLoader(workspace, globalSkillsDir, builtinSkillsDir, nil, false),
memory: NewMemoryStore(workspace),
}
}

View file

@ -327,11 +327,25 @@ func registerSharedTools(
cfg.Tools.Skills.SearchCache.MaxSize,
time.Duration(cfg.Tools.Skills.SearchCache.TTLSeconds)*time.Second,
)
agent.Tools.Register(tools.NewFindSkillsTool(registryMgr, searchCache, cfg.Tools.Skills.Whitelist, cfg.Tools.Skills.WhitelistEnabled))
agent.Tools.Register(
tools.NewFindSkillsTool(
registryMgr,
searchCache,
cfg.Tools.Skills.Whitelist,
cfg.Tools.Skills.WhitelistEnabled,
),
)
}
if install_skills_enable {
agent.Tools.Register(tools.NewInstallSkillTool(registryMgr, agent.Workspace, cfg.Tools.Skills.Whitelist, cfg.Tools.Skills.WhitelistEnabled))
agent.Tools.Register(
tools.NewInstallSkillTool(
registryMgr,
agent.Workspace,
cfg.Tools.Skills.Whitelist,
cfg.Tools.Skills.WhitelistEnabled,
),
)
}
}

View file

@ -893,7 +893,7 @@ type ToolsConfig struct {
MediaCleanup MediaCleanupConfig `json:"media_cleanup" yaml:"-"`
Whitelist FlexibleStringSlice `json:"whitelist,omitempty" yaml:"-" env:"PICOCLAW_TOOLS_WHITELIST"`
WhitelistEnabled bool `json:"whitelist_enabled,omitempty" yaml:"-" env:"PICOCLAW_TOOLS_WHITELIST_ENABLED"`
MCP MCPConfig `json:"mcp" yaml:"-""`
MCP MCPConfig `json:"mcp" yaml:"-"`
AppendFile ToolConfig `json:"append_file" yaml:"-" envPrefix:"PICOCLAW_TOOLS_APPEND_FILE_"`
EditFile ToolConfig `json:"edit_file" yaml:"-" envPrefix:"PICOCLAW_TOOLS_EDIT_FILE_"`
FindSkills ToolConfig `json:"find_skills" yaml:"-" envPrefix:"PICOCLAW_TOOLS_FIND_SKILLS_"`

View file

@ -34,8 +34,9 @@ func TestJSONUnmarshalPrivateFields(t *testing.T) {
if s.PublicField != "pub" {
t.Errorf("PublicField = %q, want 'pub'", s.PublicField)
}
// Private fields cannot be unmarshaled from JSON
if s.privateField != "" {
t.Errorf("privateField = %q, want empty because unexported fields are ignored", s.privateField)
t.Errorf("privateField = %q, want empty string (private fields are not unmarshaled)", s.privateField)
}
}

View file

@ -216,7 +216,6 @@ func Run(debug bool, homePath, configPath string, allowEmptyStartup bool) error
})
}
fmt.Printf("✓ Gateway started on %s:%d\n", cfg.Gateway.Host, cfg.Gateway.Port)
fmt.Println("Press Ctrl+C to stop")

View file

@ -2,15 +2,28 @@ package health
import (
"context"
"crypto/subtle"
"encoding/json"
"fmt"
"maps"
"net/http"
"os"
"sync"
"time"
"github.com/sipeed/picoclaw/pkg/logger"
)
// ChatRequest is the JSON body for POST /chat.
type ChatRequest struct {
Message string `json:"message"`
SessionID string `json:"session_id,omitempty"`
}
// ChatResponse is the JSON response from POST /chat.
type ChatResponse struct {
Response string `json:"response"`
}
type Server struct {
server *http.Server
mu sync.RWMutex
@ -23,7 +36,6 @@ type Server struct {
apiKey string
}
type Check struct {
Name string `json:"name"`
Status string `json:"status"`
@ -35,6 +47,7 @@ type StatusResponse struct {
Status string `json:"status"`
Uptime string `json:"uptime"`
Checks map[string]Check `json:"checks,omitempty"`
Pid int `json:"pid"`
}
func NewServer(host string, port int, token string) *Server {
@ -51,13 +64,13 @@ func NewServer(host string, port int, token string) *Server {
mux.HandleFunc("/reload", s.reloadHandler)
mux.HandleFunc("/chat", s.chatHandler)
addr := fmt.Sprintf("%s:%d", host, port)
s.server = &http.Server{
Addr: addr,
Handler: mux,
ReadTimeout: 5 * time.Second,
WriteTimeout: 5 * time.Second,
ReadTimeout: 10 * time.Second,
// WriteTimeout must be long enough for LLM inference; 5 min is generous.
WriteTimeout: 5 * time.Minute,
}
return s
@ -121,27 +134,44 @@ func (s *Server) SetReloadFunc(fn func() error) {
s.reloadFunc = fn
}
func (s *Server) reloadHandler(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusMethodNotAllowed)
json.NewEncoder(w).Encode(map[string]string{"error": "method not allowed, use POST"})
return
// SetChatFunc sets the callback that processes /chat requests.
// fn receives the user message and an optional session ID and must return the
// agent's reply (or an error). It is called synchronously inside the HTTP
// handler, so the write timeout on the server governs the maximum duration.
func (s *Server) SetChatFunc(fn func(ctx context.Context, message, sessionID string) (string, error)) {
s.mu.Lock()
defer s.mu.Unlock()
s.chatFunc = fn
}
// Token check
s.mu.RLock()
requiredToken := s.authToken
s.mu.RUnlock()
// SetAPIKey sets the expected X-API-Key header value.
func (s *Server) SetAPIKey(key string) {
s.mu.Lock()
defer s.mu.Unlock()
s.apiKey = key
}
if requiredToken != "" {
given := extractBearerToken(r.Header.Get("Authorization"))
if given == "" || subtle.ConstantTimeCompare([]byte(given), []byte(requiredToken)) != 1 {
func (s *Server) verifyAPIKey(r *http.Request) bool {
s.mu.RLock()
defer s.mu.RUnlock()
if s.apiKey == "" {
return true
}
return r.Header.Get("X-API-Key") == s.apiKey
}
func (s *Server) reloadHandler(w http.ResponseWriter, r *http.Request) {
if !s.verifyAPIKey(r) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusUnauthorized)
json.NewEncoder(w).Encode(map[string]string{"error": "unauthorized"})
return
}
if r.Method != http.MethodPost {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusMethodNotAllowed)
json.NewEncoder(w).Encode(map[string]string{"error": "method not allowed, use POST"})
return
}
s.mu.Lock()
@ -175,6 +205,7 @@ func (s *Server) healthHandler(w http.ResponseWriter, r *http.Request) {
resp := StatusResponse{
Status: "ok",
Uptime: uptime.String(),
Pid: os.Getpid(),
}
json.NewEncoder(w).Encode(resp)
@ -218,20 +249,72 @@ func (s *Server) readyHandler(w http.ResponseWriter, r *http.Request) {
})
}
// HandlerMux is the interface for registering HTTP handlers, used by
// RegisterOnMux so that callers can pass any mux implementation
// (e.g. *http.ServeMux or a custom dynamic mux).
type HandlerMux interface {
Handle(pattern string, handler http.Handler)
HandleFunc(pattern string, handler func(http.ResponseWriter, *http.Request))
}
// RegisterOnMux registers /health, /ready and /reload handlers onto the given mux.
// This allows the health endpoints to be served by a shared HTTP server.
func (s *Server) RegisterOnMux(mux HandlerMux) {
// RegisterOnMux registers /health, /ready, /reload and /chat handlers onto the
// given mux. This allows the health endpoints to be served by a shared HTTP server.
func (s *Server) RegisterOnMux(mux *http.ServeMux) {
mux.HandleFunc("/health", s.healthHandler)
mux.HandleFunc("/ready", s.readyHandler)
mux.HandleFunc("/reload", s.reloadHandler)
mux.HandleFunc("/chat", s.chatHandler)
mux.HandleFunc("/v1/chat/completions", func(w http.ResponseWriter, r *http.Request) {
logger.Error("GATEWAY IS HITTING ITSELF FOR LLM CALLS!")
http.Error(w, "GATEWAY LOOP DETECTION", http.StatusLoopDetected)
})
}
// chatHandler handles POST /chat — a synchronous HTTP chat API.
// Request body: {"message": "...", "session_id": "..." (optional)}
// Response body: {"response": "..."}
func (s *Server) chatHandler(w http.ResponseWriter, r *http.Request) {
if !s.verifyAPIKey(r) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusUnauthorized)
json.NewEncoder(w).Encode(map[string]string{"error": "unauthorized"})
return
}
if r.Method != http.MethodPost {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusMethodNotAllowed)
json.NewEncoder(w).Encode(map[string]string{"error": "method not allowed, use POST"})
return
}
s.mu.RLock()
chatFunc := s.chatFunc
s.mu.RUnlock()
if chatFunc == nil {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusServiceUnavailable)
json.NewEncoder(w).Encode(map[string]string{"error": "chat not configured"})
return
}
var req ChatRequest
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusBadRequest)
json.NewEncoder(w).Encode(map[string]string{"error": "invalid JSON: " + err.Error()})
return
}
if req.Message == "" {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusBadRequest)
json.NewEncoder(w).Encode(map[string]string{"error": "message field is required"})
return
}
reply, err := chatFunc(r.Context(), req.Message, req.SessionID)
if err != nil {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusInternalServerError)
json.NewEncoder(w).Encode(map[string]string{"error": err.Error()})
return
}
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
json.NewEncoder(w).Encode(ChatResponse{Response: reply})
}
func statusString(ok bool) string {
@ -240,102 +323,3 @@ func statusString(ok bool) string {
}
return "fail"
}
// extractBearerToken returns the token from an "Authorization: Bearer <t>" header,
// or the empty string if the header is missing or malformed.
func extractBearerToken(header string) string {
const prefix = "Bearer "
if len(header) < len(prefix) {
return ""
}
if header[:len(prefix)] != prefix {
return ""
}
return header[len(prefix):]
}
// SetChatFunc sets the callback that processes /chat requests.
func (s *Server) SetChatFunc(fn func(ctx context.Context, message, sessionID string) (string, error)) {
s.mu.Lock()
defer s.mu.Unlock()
s.chatFunc = fn
}
// SetAPIKey sets the expected X-API-Key header value.
func (s *Server) SetAPIKey(key string) {
s.mu.Lock()
defer s.mu.Unlock()
s.apiKey = key
}
func (s *Server) verifyAPIKey(r *http.Request) bool {
s.mu.RLock()
defer s.mu.RUnlock()
if s.apiKey == "" {
true
}
return r.Header.Get("X-API-Key") == s.apiKey
}
// ChatRequest is the JSON body for POST /chat.
type ChatRequest struct {
Message string `json:"message"`
SessionID string `json:"session_id,omitempty"`
}
// ChatResponse is the JSON response from POST /chat.
type ChatResponse struct {
Response string `json:"response"`
}
func (s *Server) chatHandler(w http.ResponseWriter, r *http.Request) {
if !s.verifyAPIKey(r) {
tent-Type", "application/json")
authorized)
.NewEncoder(w).Encode(map[string]string{"error": "unauthorized"})
}
if r.Method != http.MethodPost {
tent-Type", "application/json")
otAllowed)
.NewEncoder(w).Encode(map[string]string{"error": "method not allowed, use POST"})
}
s.mu.RLock()
chatFunc := s.chatFunc
s.mu.RUnlock()
if chatFunc == nil {
tent-Type", "application/json")
available)
.NewEncoder(w).Encode(map[string]string{"error": "chat not configured"})
}
var req ChatRequest
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
tent-Type", "application/json")
uest)
.NewEncoder(w).Encode(map[string]string{"error": "invalid JSON: " + err.Error()})
}
if req.Message == "" {
tent-Type", "application/json")
uest)
.NewEncoder(w).Encode(map[string]string{"error": "message field is required"})
}
reply, err := chatFunc(r.Context(), req.Message, req.SessionID)
if err != nil {
tent-Type", "application/json")
ternalServerError)
.NewEncoder(w).Encode(map[string]string{"error": err.Error()})
}
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
json.NewEncoder(w).Encode(ChatResponse{Response: reply})
}

View file

@ -240,6 +240,36 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err
cfg.ExtraBody,
), modelID, nil
case "nvidia":
apiBase := cfg.APIBase
if apiBase == "" {
apiBase = getDefaultAPIBase(protocol)
}
p := NewHTTPProviderWithMaxTokensFieldAndRequestTimeout(
cfg.APIKey(),
apiBase,
cfg.Proxy,
cfg.MaxTokensField,
cfg.RequestTimeout,
cfg.ExtraBody,
)
// NVIDIA sometimes prefers api-key header or has issues with Bearer in some environments
p.SetUseAzureHeaders(false) // NVIDIA main gateway prefers standard Bearer headers; api-key causes 404s
return p, "nvidia/" + modelID, nil
case "azure-ai", "azure-foundry":
// Azure AI Foundry / Studio compatible with OpenAI API format,
// but using api-key header instead of Authorization: Bearer.
if cfg.APIKey() == "" && cfg.APIBase == "" {
return nil, "", fmt.Errorf("api_key or api_base is required for protocol %q", protocol)
}
return NewAzureAIProvider(
cfg.APIKey(),
cfg.APIBase,
cfg.Proxy,
cfg.RequestTimeout,
), modelID, nil
case "minimax":
// Minimax requires reasoning_split: true in the request body
if cfg.APIKey() == "" && cfg.APIBase == "" {

View file

@ -91,4 +91,3 @@ func (p *HTTPProvider) SetUseAzureHeaders(use bool) {
func (p *HTTPProvider) SupportsNativeSearch() bool {
return p.delegate.SupportsNativeSearch()
}

View file

@ -11,8 +11,10 @@ import (
"net/http"
"net/url"
"strings"
"sync"
"time"
"github.com/sipeed/picoclaw/pkg/providers/common"
"github.com/sipeed/picoclaw/pkg/providers/protocoltypes"
)
@ -37,8 +39,12 @@ type Provider struct {
httpClient *http.Client
extraBody map[string]any // Additional fields to inject into request body
userAgent string
useAzureHeaders bool // Use api-key header instead of Authorization: Bearer
mu sync.RWMutex // Protect useAzureHeaders
}
type Option func(*Provider)
const defaultRequestTimeout = common.DefaultRequestTimeout
@ -90,6 +96,19 @@ func WithExtraBody(extraBody map[string]any) Option {
}
}
func WithAzureHeaders(use bool) Option {
return func(p *Provider) {
p.useAzureHeaders = use
}
}
func (p *Provider) SetUseAzureHeaders(use bool) {
p.mu.Lock()
defer p.mu.Unlock()
p.useAzureHeaders = use
}
func NewProvider(apiKey, apiBase, proxy string, opts ...Option) *Provider {
p := &Provider{
apiKey: apiKey,
@ -459,7 +478,7 @@ func isNativeSearchHost(apiBase string) bool {
return false
}
host := u.Hostname()
return host == "api.openai.com" || strings.HasSuffix(host, ".openai.azure.com")
return host == "api.openai.com"
}
// supportsPromptCacheKey reports whether the given API base is known to
@ -472,5 +491,5 @@ func supportsPromptCacheKey(apiBase string) bool {
return false
}
host := u.Hostname()
return host == "api.openai.com" || strings.HasSuffix(host, ".openai.azure.com")
return host == "api.openai.com"
}

View file

@ -90,7 +90,13 @@ func (sl *SkillsLoader) SkillRoots() []string {
return out
}
func NewSkillsLoader(workspace string, globalSkills string, builtinSkills string, whitelist []string, whitelistEnabled bool) *SkillsLoader {
func NewSkillsLoader(
workspace string,
globalSkills string,
builtinSkills string,
whitelist []string,
whitelistEnabled bool,
) *SkillsLoader {
return &SkillsLoader{
workspace: workspace,
workspaceSkills: filepath.Join(workspace, "skills"),

View file

@ -417,6 +417,7 @@ func TestGetSkillMetadata_IgnoresHTMLCommentBlocks(t *testing.T) {
assert.Equal(t, "biomed-skill", meta.Name)
assert.Equal(t, "Summarize biomedical papers.", meta.Description)
}
func TestListSkillsWithWhitelist(t *testing.T) {
tmp := t.TempDir()
ws := filepath.Join(tmp, "workspace")

View file

@ -26,7 +26,12 @@ type InstallSkillTool struct {
// NewInstallSkillTool creates a new InstallSkillTool.
// registryMgr is the shared registry manager (same instance as FindSkillsTool).
// workspace is the root workspace directory; skills install to {workspace}/skills/{slug}/.
func NewInstallSkillTool(registryMgr *skills.RegistryManager, workspace string, whitelist []string, whitelistEnabled bool) *InstallSkillTool {
func NewInstallSkillTool(
registryMgr *skills.RegistryManager,
workspace string,
whitelist []string,
whitelistEnabled bool,
) *InstallSkillTool {
return &InstallSkillTool{
registryMgr: registryMgr,
workspace: workspace,

View file

@ -102,6 +102,7 @@ func TestInstallSkillToolMissingRegistry(t *testing.T) {
assert.True(t, result.IsError)
assert.Contains(t, result.ForLLM, "invalid registry")
}
func TestInstallSkillToolWhitelist(t *testing.T) {
workspace := t.TempDir()
rm := skills.NewRegistryManager()

View file

@ -19,7 +19,12 @@ type FindSkillsTool struct {
// NewFindSkillsTool creates a new FindSkillsTool.
// registryMgr is the shared registry manager (built from config in createToolRegistry).
// cache is the search cache for deduplicating similar queries.
func NewFindSkillsTool(registryMgr *skills.RegistryManager, cache *skills.SearchCache, whitelist []string, enabled bool) *FindSkillsTool {
func NewFindSkillsTool(
registryMgr *skills.RegistryManager,
cache *skills.SearchCache,
whitelist []string,
enabled bool,
) *FindSkillsTool {
return &FindSkillsTool{
registryMgr: registryMgr,
cache: cache,
@ -98,7 +103,6 @@ func (t *FindSkillsTool) Execute(ctx context.Context, args map[string]any) *Tool
results = filtered
}
// Cache the results.
if t.cache != nil && len(results) > 0 {
t.cache.Put(query, results)

View file

@ -106,10 +106,14 @@ build-dev-picoclaw:
@mkdir -p "$$(dirname "$(PICOCLAW_BINARY)")"
@$(GO) build $(GOFLAGS) -ldflags "$(LDFLAGS)" -o "$(PICOCLAW_BINARY)" ../cmd/picoclaw
# Run all tests
test:
cd $(BACKEND_DIR) && ${WEB_GO} test ./...
cd $(FRONTEND_DIR) && pnpm lint
@if command -v pnpm >/dev/null 2>&1; then \
cd $(FRONTEND_DIR) && pnpm lint; \
else \
echo "pnpm not found, skipping frontend linting"; \
fi
# Lint and format
lint:

View file

@ -130,8 +130,12 @@ func (h *Handler) handleAddModel(w http.ResponseWriter, r *http.Request) {
return
}
if mc.APIKey != "" {
mc.ModelConfig.SetAPIKey(mc.APIKey)
apiKey := mc.APIKey
if apiKey == "" {
apiKey = mc.ModelConfig.APIKey()
}
if apiKey != "" {
mc.ModelConfig.SetAPIKey(apiKey)
}
cfg, err := config.LoadConfig(h.configPath)
@ -201,13 +205,15 @@ func (h *Handler) handleUpdateModel(w http.ResponseWriter, r *http.Request) {
return
}
// Preserve the existing API key when the caller omits it (empty string).
// This lets the UI update api_base / proxy without clearing the stored secret.
if mc.APIKey == "" {
mc.ModelConfig.SetAPIKey(cfg.ModelList[idx].APIKey())
} else {
mc.ModelConfig.SetAPIKey(mc.APIKey)
apiKey := mc.APIKey
if apiKey == "" {
apiKey = mc.ModelConfig.APIKey()
}
if apiKey == "" {
apiKey = cfg.ModelList[idx].APIKey()
}
mc.ModelConfig.SetAPIKey(apiKey)
// Preserve existing ExtraBody when omitted (nil), but clear it when
// the frontend sends an empty object {} to indicate the field should
// be removed.

View file

@ -507,6 +507,8 @@ func newSkillsLoader(workspace string) *skills.SkillsLoader {
workspace,
filepath.Join(globalConfigDir(), "skills"),
builtinSkillsDir(),
nil,
false,
)
}