Merge branch 'main' into version
This commit is contained in:
commit
36f9d20de1
24 changed files with 402 additions and 76 deletions
|
|
@ -33,7 +33,7 @@ func LoadConfig() (*config.Config, error) {
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
logger.SetLevelFromString(cfg.Agents.Defaults.LogLevel)
|
logger.SetLevelFromString(cfg.Gateway.LogLevel)
|
||||||
return cfg, nil
|
return cfg, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1,7 +1,6 @@
|
||||||
{
|
{
|
||||||
"agents": {
|
"agents": {
|
||||||
"defaults": {
|
"defaults": {
|
||||||
"log_level": "fatal",
|
|
||||||
"workspace": "~/.picoclaw/workspace",
|
"workspace": "~/.picoclaw/workspace",
|
||||||
"restrict_to_workspace": true,
|
"restrict_to_workspace": true,
|
||||||
"model_name": "gpt-5.4",
|
"model_name": "gpt-5.4",
|
||||||
|
|
@ -560,8 +559,10 @@
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
"gateway": {
|
"gateway": {
|
||||||
|
"_comment": "Default log level is set to 'fatal'. Other available options are 'debug', 'info', 'warn' and 'error'.",
|
||||||
"host": "127.0.0.1",
|
"host": "127.0.0.1",
|
||||||
"port": 18790,
|
"port": 18790,
|
||||||
"hot_reload": false
|
"hot_reload": false,
|
||||||
|
"log_level": "fatal"
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -398,6 +398,7 @@ func (c *DiscordChannel) handleMessage(s *discordgo.Session, m *discordgo.Messag
|
||||||
ref, err := store.Store(localPath, media.MediaMeta{
|
ref, err := store.Store(localPath, media.MediaMeta{
|
||||||
Filename: filename,
|
Filename: filename,
|
||||||
Source: "discord",
|
Source: "discord",
|
||||||
|
CleanupPolicy: media.CleanupPolicyDeleteOnCleanup,
|
||||||
}, scope)
|
}, scope)
|
||||||
if err == nil {
|
if err == nil {
|
||||||
return ref
|
return ref
|
||||||
|
|
|
||||||
|
|
@ -727,6 +727,7 @@ func (c *FeishuChannel) downloadResource(
|
||||||
ref, err := store.Store(localPath, media.MediaMeta{
|
ref, err := store.Store(localPath, media.MediaMeta{
|
||||||
Filename: filename,
|
Filename: filename,
|
||||||
Source: "feishu",
|
Source: "feishu",
|
||||||
|
CleanupPolicy: media.CleanupPolicyDeleteOnCleanup,
|
||||||
}, scope)
|
}, scope)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
logger.ErrorCF("feishu", "Failed to store downloaded resource", map[string]any{
|
logger.ErrorCF("feishu", "Failed to store downloaded resource", map[string]any{
|
||||||
|
|
|
||||||
|
|
@ -303,6 +303,7 @@ func (c *LINEChannel) processEvent(event lineEvent) {
|
||||||
ref, err := store.Store(localPath, media.MediaMeta{
|
ref, err := store.Store(localPath, media.MediaMeta{
|
||||||
Filename: filename,
|
Filename: filename,
|
||||||
Source: "line",
|
Source: "line",
|
||||||
|
CleanupPolicy: media.CleanupPolicyDeleteOnCleanup,
|
||||||
}, scope)
|
}, scope)
|
||||||
if err == nil {
|
if err == nil {
|
||||||
return ref
|
return ref
|
||||||
|
|
|
||||||
|
|
@ -692,6 +692,9 @@ func (c *MatrixChannel) extractInboundMedia(
|
||||||
|
|
||||||
func (c *MatrixChannel) storeMedia(localPath string, meta media.MediaMeta, scope string) string {
|
func (c *MatrixChannel) storeMedia(localPath string, meta media.MediaMeta, scope string) string {
|
||||||
if store := c.GetMediaStore(); store != nil {
|
if store := c.GetMediaStore(); store != nil {
|
||||||
|
if meta.CleanupPolicy == "" {
|
||||||
|
meta.CleanupPolicy = media.CleanupPolicyDeleteOnCleanup
|
||||||
|
}
|
||||||
ref, err := store.Store(localPath, meta, scope)
|
ref, err := store.Store(localPath, meta, scope)
|
||||||
if err == nil {
|
if err == nil {
|
||||||
return ref
|
return ref
|
||||||
|
|
|
||||||
|
|
@ -751,6 +751,7 @@ func (c *OneBotChannel) parseMessageSegments(
|
||||||
ref, err := store.Store(localPath, media.MediaMeta{
|
ref, err := store.Store(localPath, media.MediaMeta{
|
||||||
Filename: filename,
|
Filename: filename,
|
||||||
Source: "onebot",
|
Source: "onebot",
|
||||||
|
CleanupPolicy: media.CleanupPolicyDeleteOnCleanup,
|
||||||
}, scope)
|
}, scope)
|
||||||
if err == nil {
|
if err == nil {
|
||||||
return ref
|
return ref
|
||||||
|
|
|
||||||
|
|
@ -722,6 +722,7 @@ func (c *QQChannel) extractInboundAttachments(
|
||||||
Filename: qqAttachmentFilename(attachment),
|
Filename: qqAttachmentFilename(attachment),
|
||||||
ContentType: attachment.ContentType,
|
ContentType: attachment.ContentType,
|
||||||
Source: "qq",
|
Source: "qq",
|
||||||
|
CleanupPolicy: media.CleanupPolicyDeleteOnCleanup,
|
||||||
}, scope)
|
}, scope)
|
||||||
if err == nil {
|
if err == nil {
|
||||||
return ref
|
return ref
|
||||||
|
|
|
||||||
|
|
@ -329,6 +329,7 @@ func (c *SlackChannel) handleMessageEvent(ev *slackevents.MessageEvent) {
|
||||||
ref, err := store.Store(localPath, media.MediaMeta{
|
ref, err := store.Store(localPath, media.MediaMeta{
|
||||||
Filename: filename,
|
Filename: filename,
|
||||||
Source: "slack",
|
Source: "slack",
|
||||||
|
CleanupPolicy: media.CleanupPolicyDeleteOnCleanup,
|
||||||
}, scope)
|
}, scope)
|
||||||
if err == nil {
|
if err == nil {
|
||||||
return ref
|
return ref
|
||||||
|
|
|
||||||
|
|
@ -563,6 +563,7 @@ func (c *TelegramChannel) handleMessage(ctx context.Context, message *telego.Mes
|
||||||
ref, err := store.Store(localPath, media.MediaMeta{
|
ref, err := store.Store(localPath, media.MediaMeta{
|
||||||
Filename: filename,
|
Filename: filename,
|
||||||
Source: "telegram",
|
Source: "telegram",
|
||||||
|
CleanupPolicy: media.CleanupPolicyDeleteOnCleanup,
|
||||||
}, scope)
|
}, scope)
|
||||||
if err == nil {
|
if err == nil {
|
||||||
return ref
|
return ref
|
||||||
|
|
|
||||||
|
|
@ -1220,6 +1220,7 @@ func (c *WeComAIBotWSChannel) storeWSMedia(
|
||||||
ref, err := store.Store(tmpPath, media.MediaMeta{
|
ref, err := store.Store(tmpPath, media.MediaMeta{
|
||||||
Filename: msgID + ext,
|
Filename: msgID + ext,
|
||||||
Source: "wecom_aibot",
|
Source: "wecom_aibot",
|
||||||
|
CleanupPolicy: media.CleanupPolicyDeleteOnCleanup,
|
||||||
}, scope)
|
}, scope)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
os.Remove(tmpPath)
|
os.Remove(tmpPath)
|
||||||
|
|
|
||||||
|
|
@ -294,6 +294,7 @@ func (c *WeixinChannel) storeInboundBytes(
|
||||||
Filename: filename,
|
Filename: filename,
|
||||||
ContentType: contentType,
|
ContentType: contentType,
|
||||||
Source: "weixin",
|
Source: "weixin",
|
||||||
|
CleanupPolicy: media.CleanupPolicyDeleteOnCleanup,
|
||||||
}, basechannels.BuildMediaScope("weixin", chatID, messageID))
|
}, basechannels.BuildMediaScope("weixin", chatID, messageID))
|
||||||
if err != nil {
|
if err != nil {
|
||||||
os.Remove(tmpPath)
|
os.Remove(tmpPath)
|
||||||
|
|
|
||||||
|
|
@ -299,7 +299,6 @@ type AgentDefaults struct {
|
||||||
SteeringMode string `json:"steering_mode,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_STEERING_MODE"` // "one-at-a-time" (default) or "all"
|
SteeringMode string `json:"steering_mode,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_STEERING_MODE"` // "one-at-a-time" (default) or "all"
|
||||||
SubTurn SubTurnConfig `json:"subturn" envPrefix:"PICOCLAW_AGENTS_DEFAULTS_SUBTURN_"`
|
SubTurn SubTurnConfig `json:"subturn" envPrefix:"PICOCLAW_AGENTS_DEFAULTS_SUBTURN_"`
|
||||||
ToolFeedback ToolFeedbackConfig `json:"tool_feedback,omitempty"`
|
ToolFeedback ToolFeedbackConfig `json:"tool_feedback,omitempty"`
|
||||||
LogLevel string `json:"log_level,omitempty" env:"PICOCLAW_LOG_LEVEL"`
|
|
||||||
}
|
}
|
||||||
|
|
||||||
const (
|
const (
|
||||||
|
|
@ -980,6 +979,7 @@ type GatewayConfig struct {
|
||||||
Host string `json:"host" env:"PICOCLAW_GATEWAY_HOST"`
|
Host string `json:"host" env:"PICOCLAW_GATEWAY_HOST"`
|
||||||
Port int `json:"port" env:"PICOCLAW_GATEWAY_PORT"`
|
Port int `json:"port" env:"PICOCLAW_GATEWAY_PORT"`
|
||||||
HotReload bool `json:"hot_reload" env:"PICOCLAW_GATEWAY_HOT_RELOAD"`
|
HotReload bool `json:"hot_reload" env:"PICOCLAW_GATEWAY_HOT_RELOAD"`
|
||||||
|
LogLevel string `json:"log_level,omitempty" env:"PICOCLAW_LOG_LEVEL"`
|
||||||
}
|
}
|
||||||
|
|
||||||
type ToolDiscoveryConfig struct {
|
type ToolDiscoveryConfig struct {
|
||||||
|
|
@ -1323,6 +1323,9 @@ func LoadConfig(path string) (*Config, error) {
|
||||||
if e := json.Unmarshal(data, &versionInfo); e != nil {
|
if e := json.Unmarshal(data, &versionInfo); e != nil {
|
||||||
return nil, fmt.Errorf("failed to detect config version: %w", e)
|
return nil, fmt.Errorf("failed to detect config version: %w", e)
|
||||||
}
|
}
|
||||||
|
if len(data) <= 10 {
|
||||||
|
return DefaultConfig().WithSecurity(&SecurityConfig{}), nil
|
||||||
|
}
|
||||||
|
|
||||||
// Load config based on detected version
|
// Load config based on detected version
|
||||||
var cfg *Config
|
var cfg *Config
|
||||||
|
|
|
||||||
|
|
@ -461,8 +461,8 @@ func TestDefaultConfig_HooksDefaults(t *testing.T) {
|
||||||
|
|
||||||
func TestDefaultConfig_LogLevel(t *testing.T) {
|
func TestDefaultConfig_LogLevel(t *testing.T) {
|
||||||
cfg := DefaultConfig()
|
cfg := DefaultConfig()
|
||||||
if cfg.Agents.Defaults.LogLevel != "fatal" {
|
if cfg.Gateway.LogLevel != "fatal" {
|
||||||
t.Errorf("LogLevel = %q, want \"fatal\"", cfg.Agents.Defaults.LogLevel)
|
t.Errorf("LogLevel = %q, want \"fatal\"", cfg.Gateway.LogLevel)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -1162,7 +1162,7 @@ func TestLoadConfig_UsesPassphraseProvider(t *testing.T) {
|
||||||
func TestConfigParsesLogLevel(t *testing.T) {
|
func TestConfigParsesLogLevel(t *testing.T) {
|
||||||
dir := t.TempDir()
|
dir := t.TempDir()
|
||||||
cfgPath := filepath.Join(dir, "config.json")
|
cfgPath := filepath.Join(dir, "config.json")
|
||||||
data := `{"version":1,"agents":{"defaults":{"log_level":"debug"}}}`
|
data := `{"version":1,"gateway":{"log_level":"debug"}}`
|
||||||
if err := os.WriteFile(cfgPath, []byte(data), 0o600); err != nil {
|
if err := os.WriteFile(cfgPath, []byte(data), 0o600); err != nil {
|
||||||
t.Fatalf("setup: %v", err)
|
t.Fatalf("setup: %v", err)
|
||||||
}
|
}
|
||||||
|
|
@ -1171,15 +1171,15 @@ func TestConfigParsesLogLevel(t *testing.T) {
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatalf("LoadConfig: %v", err)
|
t.Fatalf("LoadConfig: %v", err)
|
||||||
}
|
}
|
||||||
if cfg.Agents.Defaults.LogLevel != "debug" {
|
if cfg.Gateway.LogLevel != "debug" {
|
||||||
t.Errorf("LogLevel = %q, want \"debug\"", cfg.Agents.Defaults.LogLevel)
|
t.Errorf("LogLevel = %q, want \"debug\"", cfg.Gateway.LogLevel)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestConfigLogLevelEmpty(t *testing.T) {
|
func TestConfigLogLevelEmpty(t *testing.T) {
|
||||||
dir := t.TempDir()
|
dir := t.TempDir()
|
||||||
cfgPath := filepath.Join(dir, "config.json")
|
cfgPath := filepath.Join(dir, "config.json")
|
||||||
data := `{}`
|
data := `{"version":1}`
|
||||||
if err := os.WriteFile(cfgPath, []byte(data), 0o600); err != nil {
|
if err := os.WriteFile(cfgPath, []byte(data), 0o600); err != nil {
|
||||||
t.Fatalf("setup: %v", err)
|
t.Fatalf("setup: %v", err)
|
||||||
}
|
}
|
||||||
|
|
@ -1189,7 +1189,7 @@ func TestConfigLogLevelEmpty(t *testing.T) {
|
||||||
t.Fatalf("LoadConfig: %v", err)
|
t.Fatalf("LoadConfig: %v", err)
|
||||||
}
|
}
|
||||||
// When config omits log_level, the DefaultConfig value ("fatal") is preserved.
|
// When config omits log_level, the DefaultConfig value ("fatal") is preserved.
|
||||||
if cfg.Agents.Defaults.LogLevel != "fatal" {
|
if cfg.Gateway.LogLevel != "fatal" {
|
||||||
t.Errorf("LogLevel = %q, want \"fatal\"", cfg.Agents.Defaults.LogLevel)
|
t.Errorf("LogLevel = %q, want \"fatal\"", cfg.Gateway.LogLevel)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -29,7 +29,6 @@ func DefaultConfig() *Config {
|
||||||
Version: CurrentVersion,
|
Version: CurrentVersion,
|
||||||
Agents: AgentsConfig{
|
Agents: AgentsConfig{
|
||||||
Defaults: AgentDefaults{
|
Defaults: AgentDefaults{
|
||||||
LogLevel: "fatal",
|
|
||||||
Workspace: workspacePath,
|
Workspace: workspacePath,
|
||||||
RestrictToWorkspace: true,
|
RestrictToWorkspace: true,
|
||||||
Provider: "",
|
Provider: "",
|
||||||
|
|
@ -375,6 +374,7 @@ func DefaultConfig() *Config {
|
||||||
Host: "127.0.0.1",
|
Host: "127.0.0.1",
|
||||||
Port: 18790,
|
Port: 18790,
|
||||||
HotReload: false,
|
HotReload: false,
|
||||||
|
LogLevel: "fatal",
|
||||||
},
|
},
|
||||||
Tools: ToolsConfig{
|
Tools: ToolsConfig{
|
||||||
MediaCleanup: MediaCleanupConfig{
|
MediaCleanup: MediaCleanupConfig{
|
||||||
|
|
|
||||||
|
|
@ -85,7 +85,7 @@ func Run(debug bool, configPath string, allowEmptyStartup bool) error {
|
||||||
return fmt.Errorf("error loading config: %w", err)
|
return fmt.Errorf("error loading config: %w", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
logger.SetLevelFromString(cfg.Agents.Defaults.LogLevel)
|
logger.SetLevelFromString(cfg.Gateway.LogLevel)
|
||||||
|
|
||||||
if debug {
|
if debug {
|
||||||
logger.SetLevel(logger.DEBUG)
|
logger.SetLevel(logger.DEBUG)
|
||||||
|
|
|
||||||
|
|
@ -11,11 +11,25 @@ import (
|
||||||
"github.com/sipeed/picoclaw/pkg/logger"
|
"github.com/sipeed/picoclaw/pkg/logger"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
// CleanupPolicy controls how the MediaStore treats the underlying file when
|
||||||
|
// a ref is released or expires.
|
||||||
|
type CleanupPolicy string
|
||||||
|
|
||||||
|
const (
|
||||||
|
// CleanupPolicyDeleteOnCleanup means the file is store-managed and may be
|
||||||
|
// deleted once the final ref for that path is gone.
|
||||||
|
CleanupPolicyDeleteOnCleanup CleanupPolicy = "delete_on_cleanup"
|
||||||
|
// CleanupPolicyForgetOnly means the store should only drop ref mappings and
|
||||||
|
// must never delete the underlying file.
|
||||||
|
CleanupPolicyForgetOnly CleanupPolicy = "forget_only"
|
||||||
|
)
|
||||||
|
|
||||||
// MediaMeta holds metadata about a stored media file.
|
// MediaMeta holds metadata about a stored media file.
|
||||||
type MediaMeta struct {
|
type MediaMeta struct {
|
||||||
Filename string
|
Filename string
|
||||||
ContentType string
|
ContentType string
|
||||||
Source string // "telegram", "discord", "tool:image-gen", etc.
|
Source string // "telegram", "discord", "tool:image-gen", etc.
|
||||||
|
CleanupPolicy CleanupPolicy // defaults to CleanupPolicyDeleteOnCleanup
|
||||||
}
|
}
|
||||||
|
|
||||||
// MediaStore manages the lifecycle of media files associated with processing scopes.
|
// MediaStore manages the lifecycle of media files associated with processing scopes.
|
||||||
|
|
@ -23,6 +37,7 @@ type MediaStore interface {
|
||||||
// Store registers an existing local file under the given scope.
|
// Store registers an existing local file under the given scope.
|
||||||
// Returns a ref identifier (e.g. "media://<id>").
|
// Returns a ref identifier (e.g. "media://<id>").
|
||||||
// Store does not move or copy the file; it only records the mapping.
|
// Store does not move or copy the file; it only records the mapping.
|
||||||
|
// If meta.CleanupPolicy is empty, CleanupPolicyDeleteOnCleanup is assumed.
|
||||||
Store(localPath string, meta MediaMeta, scope string) (ref string, err error)
|
Store(localPath string, meta MediaMeta, scope string) (ref string, err error)
|
||||||
|
|
||||||
// Resolve returns the local file path for a given ref.
|
// Resolve returns the local file path for a given ref.
|
||||||
|
|
@ -43,6 +58,11 @@ type mediaEntry struct {
|
||||||
storedAt time.Time
|
storedAt time.Time
|
||||||
}
|
}
|
||||||
|
|
||||||
|
type pathRefState struct {
|
||||||
|
refCount int
|
||||||
|
deleteEligible bool
|
||||||
|
}
|
||||||
|
|
||||||
// MediaCleanerConfig configures the background TTL cleanup.
|
// MediaCleanerConfig configures the background TTL cleanup.
|
||||||
type MediaCleanerConfig struct {
|
type MediaCleanerConfig struct {
|
||||||
Enabled bool
|
Enabled bool
|
||||||
|
|
@ -57,6 +77,8 @@ type FileMediaStore struct {
|
||||||
refs map[string]mediaEntry
|
refs map[string]mediaEntry
|
||||||
scopeToRefs map[string]map[string]struct{}
|
scopeToRefs map[string]map[string]struct{}
|
||||||
refToScope map[string]string
|
refToScope map[string]string
|
||||||
|
refToPath map[string]string
|
||||||
|
pathStates map[string]pathRefState
|
||||||
|
|
||||||
cleanerCfg MediaCleanerConfig
|
cleanerCfg MediaCleanerConfig
|
||||||
stop chan struct{}
|
stop chan struct{}
|
||||||
|
|
@ -71,6 +93,8 @@ func NewFileMediaStore() *FileMediaStore {
|
||||||
refs: make(map[string]mediaEntry),
|
refs: make(map[string]mediaEntry),
|
||||||
scopeToRefs: make(map[string]map[string]struct{}),
|
scopeToRefs: make(map[string]map[string]struct{}),
|
||||||
refToScope: make(map[string]string),
|
refToScope: make(map[string]string),
|
||||||
|
refToPath: make(map[string]string),
|
||||||
|
pathStates: make(map[string]pathRefState),
|
||||||
nowFunc: time.Now,
|
nowFunc: time.Now,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
@ -81,6 +105,8 @@ func NewFileMediaStoreWithCleanup(cfg MediaCleanerConfig) *FileMediaStore {
|
||||||
refs: make(map[string]mediaEntry),
|
refs: make(map[string]mediaEntry),
|
||||||
scopeToRefs: make(map[string]map[string]struct{}),
|
scopeToRefs: make(map[string]map[string]struct{}),
|
||||||
refToScope: make(map[string]string),
|
refToScope: make(map[string]string),
|
||||||
|
refToPath: make(map[string]string),
|
||||||
|
pathStates: make(map[string]pathRefState),
|
||||||
cleanerCfg: cfg,
|
cleanerCfg: cfg,
|
||||||
stop: make(chan struct{}),
|
stop: make(chan struct{}),
|
||||||
nowFunc: time.Now,
|
nowFunc: time.Now,
|
||||||
|
|
@ -94,6 +120,7 @@ func (s *FileMediaStore) Store(localPath string, meta MediaMeta, scope string) (
|
||||||
}
|
}
|
||||||
|
|
||||||
ref := "media://" + uuid.New().String()
|
ref := "media://" + uuid.New().String()
|
||||||
|
meta.CleanupPolicy = normalizeCleanupPolicy(meta.CleanupPolicy)
|
||||||
|
|
||||||
s.mu.Lock()
|
s.mu.Lock()
|
||||||
defer s.mu.Unlock()
|
defer s.mu.Unlock()
|
||||||
|
|
@ -104,6 +131,18 @@ func (s *FileMediaStore) Store(localPath string, meta MediaMeta, scope string) (
|
||||||
}
|
}
|
||||||
s.scopeToRefs[scope][ref] = struct{}{}
|
s.scopeToRefs[scope][ref] = struct{}{}
|
||||||
s.refToScope[ref] = scope
|
s.refToScope[ref] = scope
|
||||||
|
s.refToPath[ref] = localPath
|
||||||
|
|
||||||
|
pathState := s.pathStates[localPath]
|
||||||
|
if pathState.refCount == 0 {
|
||||||
|
pathState.deleteEligible = meta.CleanupPolicy == CleanupPolicyDeleteOnCleanup
|
||||||
|
} else if meta.CleanupPolicy == CleanupPolicyForgetOnly {
|
||||||
|
// Be conservative: once a path is borrowed externally, never let this
|
||||||
|
// lifecycle auto-delete it even if store-managed refs also exist.
|
||||||
|
pathState.deleteEligible = false
|
||||||
|
}
|
||||||
|
pathState.refCount++
|
||||||
|
s.pathStates[localPath] = pathState
|
||||||
|
|
||||||
return ref, nil
|
return ref, nil
|
||||||
}
|
}
|
||||||
|
|
@ -134,7 +173,8 @@ func (s *FileMediaStore) ResolveWithMeta(ref string) (string, MediaMeta, error)
|
||||||
|
|
||||||
// ReleaseAll removes all files under the given scope and cleans up mappings.
|
// ReleaseAll removes all files under the given scope and cleans up mappings.
|
||||||
// Phase 1 (under lock): remove entries from maps.
|
// Phase 1 (under lock): remove entries from maps.
|
||||||
// Phase 2 (no lock): delete files from disk.
|
// Phase 2 (no lock): delete store-managed files from disk once their final
|
||||||
|
// path ref is gone.
|
||||||
func (s *FileMediaStore) ReleaseAll(scope string) error {
|
func (s *FileMediaStore) ReleaseAll(scope string) error {
|
||||||
// Phase 1: collect paths and remove from maps under lock
|
// Phase 1: collect paths and remove from maps under lock
|
||||||
var paths []string
|
var paths []string
|
||||||
|
|
@ -147,11 +187,13 @@ func (s *FileMediaStore) ReleaseAll(scope string) error {
|
||||||
}
|
}
|
||||||
|
|
||||||
for ref := range refs {
|
for ref := range refs {
|
||||||
|
fallbackPath := ""
|
||||||
if entry, exists := s.refs[ref]; exists {
|
if entry, exists := s.refs[ref]; exists {
|
||||||
paths = append(paths, entry.path)
|
fallbackPath = entry.path
|
||||||
|
}
|
||||||
|
if removablePath, shouldDelete := s.releaseRefLocked(ref, fallbackPath); shouldDelete {
|
||||||
|
paths = append(paths, removablePath)
|
||||||
}
|
}
|
||||||
delete(s.refs, ref)
|
|
||||||
delete(s.refToScope, ref)
|
|
||||||
}
|
}
|
||||||
delete(s.scopeToRefs, scope)
|
delete(s.scopeToRefs, scope)
|
||||||
s.mu.Unlock()
|
s.mu.Unlock()
|
||||||
|
|
@ -171,7 +213,7 @@ func (s *FileMediaStore) ReleaseAll(scope string) error {
|
||||||
|
|
||||||
// CleanExpired removes all entries older than MaxAge.
|
// CleanExpired removes all entries older than MaxAge.
|
||||||
// Phase 1 (under lock): identify expired entries and remove from maps.
|
// Phase 1 (under lock): identify expired entries and remove from maps.
|
||||||
// Phase 2 (no lock): delete files from disk to minimize lock contention.
|
// Phase 2 (no lock): delete store-managed files from disk to minimize lock contention.
|
||||||
func (s *FileMediaStore) CleanExpired() int {
|
func (s *FileMediaStore) CleanExpired() int {
|
||||||
if s.cleanerCfg.MaxAge <= 0 {
|
if s.cleanerCfg.MaxAge <= 0 {
|
||||||
return 0
|
return 0
|
||||||
|
|
@ -180,7 +222,7 @@ func (s *FileMediaStore) CleanExpired() int {
|
||||||
// Phase 1: collect expired entries under lock
|
// Phase 1: collect expired entries under lock
|
||||||
type expiredEntry struct {
|
type expiredEntry struct {
|
||||||
ref string
|
ref string
|
||||||
path string
|
deletePath string
|
||||||
}
|
}
|
||||||
|
|
||||||
s.mu.Lock()
|
s.mu.Lock()
|
||||||
|
|
@ -189,8 +231,6 @@ func (s *FileMediaStore) CleanExpired() int {
|
||||||
|
|
||||||
for ref, entry := range s.refs {
|
for ref, entry := range s.refs {
|
||||||
if entry.storedAt.Before(cutoff) {
|
if entry.storedAt.Before(cutoff) {
|
||||||
expired = append(expired, expiredEntry{ref: ref, path: entry.path})
|
|
||||||
|
|
||||||
if scope, ok := s.refToScope[ref]; ok {
|
if scope, ok := s.refToScope[ref]; ok {
|
||||||
if scopeRefs, ok := s.scopeToRefs[scope]; ok {
|
if scopeRefs, ok := s.scopeToRefs[scope]; ok {
|
||||||
delete(scopeRefs, ref)
|
delete(scopeRefs, ref)
|
||||||
|
|
@ -200,17 +240,23 @@ func (s *FileMediaStore) CleanExpired() int {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
delete(s.refs, ref)
|
expiredItem := expiredEntry{ref: ref}
|
||||||
delete(s.refToScope, ref)
|
if deletePath, shouldDelete := s.releaseRefLocked(ref, entry.path); shouldDelete {
|
||||||
|
expiredItem.deletePath = deletePath
|
||||||
|
}
|
||||||
|
expired = append(expired, expiredItem)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
s.mu.Unlock()
|
s.mu.Unlock()
|
||||||
|
|
||||||
// Phase 2: delete files without holding the lock
|
// Phase 2: delete files without holding the lock
|
||||||
for _, e := range expired {
|
for _, e := range expired {
|
||||||
if err := os.Remove(e.path); err != nil && !os.IsNotExist(err) {
|
if e.deletePath == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if err := os.Remove(e.deletePath); err != nil && !os.IsNotExist(err) {
|
||||||
logger.WarnCF("media", "cleanup: failed to remove file", map[string]any{
|
logger.WarnCF("media", "cleanup: failed to remove file", map[string]any{
|
||||||
"path": e.path,
|
"path": e.deletePath,
|
||||||
"error": err.Error(),
|
"error": err.Error(),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
@ -219,6 +265,45 @@ func (s *FileMediaStore) CleanExpired() int {
|
||||||
return len(expired)
|
return len(expired)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func normalizeCleanupPolicy(policy CleanupPolicy) CleanupPolicy {
|
||||||
|
switch policy {
|
||||||
|
case "", CleanupPolicyDeleteOnCleanup:
|
||||||
|
return CleanupPolicyDeleteOnCleanup
|
||||||
|
case CleanupPolicyForgetOnly:
|
||||||
|
return CleanupPolicyForgetOnly
|
||||||
|
default:
|
||||||
|
return CleanupPolicyDeleteOnCleanup
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *FileMediaStore) releaseRefLocked(ref, fallbackPath string) (string, bool) {
|
||||||
|
path := fallbackPath
|
||||||
|
if storedPath, ok := s.refToPath[ref]; ok {
|
||||||
|
path = storedPath
|
||||||
|
delete(s.refToPath, ref)
|
||||||
|
}
|
||||||
|
|
||||||
|
delete(s.refs, ref)
|
||||||
|
delete(s.refToScope, ref)
|
||||||
|
|
||||||
|
if path == "" {
|
||||||
|
return "", false
|
||||||
|
}
|
||||||
|
|
||||||
|
pathState, ok := s.pathStates[path]
|
||||||
|
if !ok {
|
||||||
|
return "", false
|
||||||
|
}
|
||||||
|
if pathState.refCount <= 1 {
|
||||||
|
delete(s.pathStates, path)
|
||||||
|
return path, pathState.deleteEligible
|
||||||
|
}
|
||||||
|
|
||||||
|
pathState.refCount--
|
||||||
|
s.pathStates[path] = pathState
|
||||||
|
return "", false
|
||||||
|
}
|
||||||
|
|
||||||
// Start begins the background cleanup goroutine if cleanup is enabled.
|
// Start begins the background cleanup goroutine if cleanup is enabled.
|
||||||
// Safe to call multiple times; only the first call starts the goroutine.
|
// Safe to call multiple times; only the first call starts the goroutine.
|
||||||
func (s *FileMediaStore) Start() {
|
func (s *FileMediaStore) Start() {
|
||||||
|
|
|
||||||
|
|
@ -77,6 +77,106 @@ func TestReleaseAll(t *testing.T) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestReleaseAllForgetOnlyKeepsFile(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
store := NewFileMediaStore()
|
||||||
|
|
||||||
|
path := createTempFile(t, dir, "workspace.txt")
|
||||||
|
ref, err := store.Store(path, MediaMeta{
|
||||||
|
Source: "test",
|
||||||
|
CleanupPolicy: CleanupPolicyForgetOnly,
|
||||||
|
}, "scope1")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Store failed: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := store.ReleaseAll("scope1"); err != nil {
|
||||||
|
t.Fatalf("ReleaseAll failed: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if _, err := store.Resolve(ref); err == nil {
|
||||||
|
t.Error("forget-only ref should be unresolvable after release")
|
||||||
|
}
|
||||||
|
if _, err := os.Stat(path); err != nil {
|
||||||
|
t.Errorf("forget-only file should remain on disk: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestReleaseAllSharedPathDeletesOnFinalRefOnly(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
store := NewFileMediaStore()
|
||||||
|
|
||||||
|
path := createTempFile(t, dir, "shared.jpg")
|
||||||
|
refA, err := store.Store(path, MediaMeta{
|
||||||
|
Source: "test",
|
||||||
|
CleanupPolicy: CleanupPolicyDeleteOnCleanup,
|
||||||
|
}, "scopeA")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Store(scopeA) failed: %v", err)
|
||||||
|
}
|
||||||
|
refB, err := store.Store(path, MediaMeta{
|
||||||
|
Source: "test",
|
||||||
|
CleanupPolicy: CleanupPolicyDeleteOnCleanup,
|
||||||
|
}, "scopeB")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Store(scopeB) failed: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := store.ReleaseAll("scopeA"); err != nil {
|
||||||
|
t.Fatalf("ReleaseAll(scopeA) failed: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if _, err := store.Resolve(refA); err == nil {
|
||||||
|
t.Error("refA should be unresolvable after ReleaseAll(scopeA)")
|
||||||
|
}
|
||||||
|
if _, err := store.Resolve(refB); err != nil {
|
||||||
|
t.Fatalf("refB should still resolve: %v", err)
|
||||||
|
}
|
||||||
|
if _, err := os.Stat(path); err != nil {
|
||||||
|
t.Errorf("shared file should remain until final ref is released: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := store.ReleaseAll("scopeB"); err != nil {
|
||||||
|
t.Fatalf("ReleaseAll(scopeB) failed: %v", err)
|
||||||
|
}
|
||||||
|
if _, err := os.Stat(path); !os.IsNotExist(err) {
|
||||||
|
t.Error("shared file should be deleted after final ref is released")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestReleaseAllMixedPoliciesKeepsFile(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
store := NewFileMediaStore()
|
||||||
|
|
||||||
|
path := createTempFile(t, dir, "shared.txt")
|
||||||
|
if _, err := store.Store(path, MediaMeta{
|
||||||
|
Source: "test",
|
||||||
|
CleanupPolicy: CleanupPolicyDeleteOnCleanup,
|
||||||
|
}, "owned"); err != nil {
|
||||||
|
t.Fatalf("Store(owned) failed: %v", err)
|
||||||
|
}
|
||||||
|
if _, err := store.Store(path, MediaMeta{
|
||||||
|
Source: "test",
|
||||||
|
CleanupPolicy: CleanupPolicyForgetOnly,
|
||||||
|
}, "borrowed"); err != nil {
|
||||||
|
t.Fatalf("Store(borrowed) failed: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := store.ReleaseAll("owned"); err != nil {
|
||||||
|
t.Fatalf("ReleaseAll(owned) failed: %v", err)
|
||||||
|
}
|
||||||
|
if _, err := os.Stat(path); err != nil {
|
||||||
|
t.Fatalf("mixed-policy file should remain after owned ref release: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := store.ReleaseAll("borrowed"); err != nil {
|
||||||
|
t.Fatalf("ReleaseAll(borrowed) failed: %v", err)
|
||||||
|
}
|
||||||
|
if _, err := os.Stat(path); err != nil {
|
||||||
|
t.Errorf("mixed-policy path should not be auto-deleted: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestMultiScopeIsolation(t *testing.T) {
|
func TestMultiScopeIsolation(t *testing.T) {
|
||||||
dir := t.TempDir()
|
dir := t.TempDir()
|
||||||
store := NewFileMediaStore()
|
store := NewFileMediaStore()
|
||||||
|
|
@ -293,6 +393,35 @@ func TestCleanExpiredRemovesOldEntries(t *testing.T) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestCleanExpiredForgetOnlyKeepsFile(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
now := time.Now()
|
||||||
|
store := newTestStoreWithCleanup(10 * time.Minute)
|
||||||
|
store.nowFunc = func() time.Time { return now.Add(-20 * time.Minute) }
|
||||||
|
|
||||||
|
path := createTempFile(t, dir, "workspace.txt")
|
||||||
|
ref, err := store.Store(path, MediaMeta{
|
||||||
|
Source: "test",
|
||||||
|
CleanupPolicy: CleanupPolicyForgetOnly,
|
||||||
|
}, "scope1")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Store failed: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
store.nowFunc = func() time.Time { return now }
|
||||||
|
removed := store.CleanExpired()
|
||||||
|
|
||||||
|
if removed != 1 {
|
||||||
|
t.Errorf("expected 1 removed, got %d", removed)
|
||||||
|
}
|
||||||
|
if _, err := store.Resolve(ref); err == nil {
|
||||||
|
t.Error("expired forget-only ref should be unresolvable")
|
||||||
|
}
|
||||||
|
if _, err := os.Stat(path); err != nil {
|
||||||
|
t.Errorf("forget-only file should remain on disk: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestCleanExpiredKeepsNonExpired(t *testing.T) {
|
func TestCleanExpiredKeepsNonExpired(t *testing.T) {
|
||||||
dir := t.TempDir()
|
dir := t.TempDir()
|
||||||
now := time.Now()
|
now := time.Now()
|
||||||
|
|
@ -346,6 +475,53 @@ func TestCleanExpiredMixedAges(t *testing.T) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestCleanExpiredSharedPathDeletesOnFinalRefOnly(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
now := time.Now()
|
||||||
|
store := newTestStoreWithCleanup(10 * time.Minute)
|
||||||
|
|
||||||
|
path := createTempFile(t, dir, "shared.jpg")
|
||||||
|
|
||||||
|
store.nowFunc = func() time.Time { return now.Add(-20 * time.Minute) }
|
||||||
|
oldRef, err := store.Store(path, MediaMeta{
|
||||||
|
Source: "test",
|
||||||
|
CleanupPolicy: CleanupPolicyDeleteOnCleanup,
|
||||||
|
}, "scope-old")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Store(old) failed: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
store.nowFunc = func() time.Time { return now }
|
||||||
|
freshRef, err := store.Store(path, MediaMeta{
|
||||||
|
Source: "test",
|
||||||
|
CleanupPolicy: CleanupPolicyDeleteOnCleanup,
|
||||||
|
}, "scope-fresh")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Store(fresh) failed: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
removed := store.CleanExpired()
|
||||||
|
if removed != 1 {
|
||||||
|
t.Errorf("expected 1 removed, got %d", removed)
|
||||||
|
}
|
||||||
|
if _, err := store.Resolve(oldRef); err == nil {
|
||||||
|
t.Error("old ref should be gone after cleanup")
|
||||||
|
}
|
||||||
|
if _, err := store.Resolve(freshRef); err != nil {
|
||||||
|
t.Fatalf("fresh ref should still resolve: %v", err)
|
||||||
|
}
|
||||||
|
if _, err := os.Stat(path); err != nil {
|
||||||
|
t.Errorf("shared file should remain while fresh ref exists: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := store.ReleaseAll("scope-fresh"); err != nil {
|
||||||
|
t.Fatalf("ReleaseAll(scope-fresh) failed: %v", err)
|
||||||
|
}
|
||||||
|
if _, err := os.Stat(path); !os.IsNotExist(err) {
|
||||||
|
t.Error("shared file should be deleted after final ref is released")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestCleanExpiredCleansEmptyScopes(t *testing.T) {
|
func TestCleanExpiredCleansEmptyScopes(t *testing.T) {
|
||||||
dir := t.TempDir()
|
dir := t.TempDir()
|
||||||
now := time.Now()
|
now := time.Now()
|
||||||
|
|
|
||||||
|
|
@ -136,6 +136,7 @@ func (t *SendFileTool) Execute(ctx context.Context, args map[string]any) *ToolRe
|
||||||
Filename: filename,
|
Filename: filename,
|
||||||
ContentType: mediaType,
|
ContentType: mediaType,
|
||||||
Source: "tool:send_file",
|
Source: "tool:send_file",
|
||||||
|
CleanupPolicy: media.CleanupPolicyForgetOnly,
|
||||||
}, scope)
|
}, scope)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return ErrorResult(fmt.Sprintf("failed to register media: %v", err))
|
return ErrorResult(fmt.Sprintf("failed to register media: %v", err))
|
||||||
|
|
|
||||||
|
|
@ -104,6 +104,14 @@ func TestSendFileTool_Success(t *testing.T) {
|
||||||
if result.Media[0][:8] != "media://" {
|
if result.Media[0][:8] != "media://" {
|
||||||
t.Errorf("expected media:// ref, got %q", result.Media[0])
|
t.Errorf("expected media:// ref, got %q", result.Media[0])
|
||||||
}
|
}
|
||||||
|
|
||||||
|
_, meta, err := store.ResolveWithMeta(result.Media[0])
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ResolveWithMeta failed: %v", err)
|
||||||
|
}
|
||||||
|
if meta.CleanupPolicy != media.CleanupPolicyForgetOnly {
|
||||||
|
t.Errorf("CleanupPolicy = %q, want %q", meta.CleanupPolicy, media.CleanupPolicyForgetOnly)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestSendFileTool_CustomFilename(t *testing.T) {
|
func TestSendFileTool_CustomFilename(t *testing.T) {
|
||||||
|
|
|
||||||
|
|
@ -169,7 +169,7 @@ func TestGatewayStartReady_LocalModelWithoutAPIKey(t *testing.T) {
|
||||||
defer cleanup()
|
defer cleanup()
|
||||||
resetModelProbeHooks(t)
|
resetModelProbeHooks(t)
|
||||||
|
|
||||||
probeOpenAICompatibleModelFunc = func(apiBase, modelID string) bool {
|
probeOpenAICompatibleModelFunc = func(apiBase, modelID, apiKey string) bool {
|
||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -206,8 +206,8 @@ func TestGatewayStartReady_LocalModelWithRunningService(t *testing.T) {
|
||||||
defer cleanup()
|
defer cleanup()
|
||||||
resetModelProbeHooks(t)
|
resetModelProbeHooks(t)
|
||||||
|
|
||||||
probeOpenAICompatibleModelFunc = func(apiBase, modelID string) bool {
|
probeOpenAICompatibleModelFunc = func(apiBase, modelID, apiKey string) bool {
|
||||||
return apiBase == "http://127.0.0.1:8000/v1" && modelID == "custom-model"
|
return apiBase == "http://127.0.0.1:8000/v1" && modelID == "custom-model" && apiKey == ""
|
||||||
}
|
}
|
||||||
|
|
||||||
cfg, err := config.LoadConfig(configPath)
|
cfg, err := config.LoadConfig(configPath)
|
||||||
|
|
@ -240,7 +240,7 @@ func TestGatewayStartReady_RemoteVLLMWithAPIKeyDoesNotProbe(t *testing.T) {
|
||||||
defer cleanup()
|
defer cleanup()
|
||||||
resetModelProbeHooks(t)
|
resetModelProbeHooks(t)
|
||||||
|
|
||||||
probeOpenAICompatibleModelFunc = func(apiBase, modelID string) bool {
|
probeOpenAICompatibleModelFunc = func(apiBase, modelID, apiKey string) bool {
|
||||||
t.Fatalf("unexpected OpenAI-compatible probe for %q (%q)", apiBase, modelID)
|
t.Fatalf("unexpected OpenAI-compatible probe for %q (%q)", apiBase, modelID)
|
||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -82,14 +82,14 @@ func probeLocalModelAvailability(m *config.ModelConfig) bool {
|
||||||
case "ollama":
|
case "ollama":
|
||||||
return probeOllamaModelFunc(apiBase, modelID)
|
return probeOllamaModelFunc(apiBase, modelID)
|
||||||
case "vllm":
|
case "vllm":
|
||||||
return probeOpenAICompatibleModelFunc(apiBase, modelID)
|
return probeOpenAICompatibleModelFunc(apiBase, modelID, m.APIKey())
|
||||||
case "github-copilot", "copilot":
|
case "github-copilot", "copilot":
|
||||||
return probeTCPServiceFunc(apiBase)
|
return probeTCPServiceFunc(apiBase)
|
||||||
case "claude-cli", "claudecli", "codex-cli", "codexcli":
|
case "claude-cli", "claudecli", "codex-cli", "codexcli":
|
||||||
return true
|
return true
|
||||||
default:
|
default:
|
||||||
if hasLocalAPIBase(apiBase) {
|
if hasLocalAPIBase(apiBase) {
|
||||||
return probeOpenAICompatibleModelFunc(apiBase, modelID)
|
return probeOpenAICompatibleModelFunc(apiBase, modelID, m.APIKey())
|
||||||
}
|
}
|
||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
|
|
@ -209,7 +209,7 @@ func probeOllamaModel(apiBase, modelID string) bool {
|
||||||
Model string `json:"model"`
|
Model string `json:"model"`
|
||||||
} `json:"models"`
|
} `json:"models"`
|
||||||
}
|
}
|
||||||
if err := getJSON(root+"/api/tags", &resp); err != nil {
|
if err := getJSON(root+"/api/tags", &resp, ""); err != nil {
|
||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -221,7 +221,7 @@ func probeOllamaModel(apiBase, modelID string) bool {
|
||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
|
|
||||||
func probeOpenAICompatibleModel(apiBase, modelID string) bool {
|
func probeOpenAICompatibleModel(apiBase, modelID, apiKey string) bool {
|
||||||
if strings.TrimSpace(apiBase) == "" {
|
if strings.TrimSpace(apiBase) == "" {
|
||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
|
|
@ -231,7 +231,7 @@ func probeOpenAICompatibleModel(apiBase, modelID string) bool {
|
||||||
ID string `json:"id"`
|
ID string `json:"id"`
|
||||||
} `json:"data"`
|
} `json:"data"`
|
||||||
}
|
}
|
||||||
if err := getJSON(strings.TrimRight(strings.TrimSpace(apiBase), "/")+"/models", &resp); err != nil {
|
if err := getJSON(strings.TrimRight(strings.TrimSpace(apiBase), "/")+"/models", &resp, apiKey); err != nil {
|
||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -243,11 +243,14 @@ func probeOpenAICompatibleModel(apiBase, modelID string) bool {
|
||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
|
|
||||||
func getJSON(rawURL string, out any) error {
|
func getJSON(rawURL string, out any, apiKey string) error {
|
||||||
req, err := http.NewRequest(http.MethodGet, rawURL, nil)
|
req, err := http.NewRequest(http.MethodGet, rawURL, nil)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
if apiKey = strings.TrimSpace(apiKey); apiKey != "" {
|
||||||
|
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||||||
|
}
|
||||||
|
|
||||||
client := &http.Client{Timeout: modelProbeTimeout}
|
client := &http.Client{Timeout: modelProbeTimeout}
|
||||||
resp, err := client.Do(req)
|
resp, err := client.Do(req)
|
||||||
|
|
|
||||||
37
web/backend/api/model_status_test.go
Normal file
37
web/backend/api/model_status_test.go
Normal file
|
|
@ -0,0 +1,37 @@
|
||||||
|
package api
|
||||||
|
|
||||||
|
import (
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/config"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestProbeLocalModelAvailability_OpenAICompatibleIncludesAPIKey(t *testing.T) {
|
||||||
|
const apiKey = "test-api-key"
|
||||||
|
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if r.URL.Path != "/v1/models" {
|
||||||
|
t.Fatalf("path = %q, want %q", r.URL.Path, "/v1/models")
|
||||||
|
}
|
||||||
|
if got := r.Header.Get("Authorization"); got != "Bearer "+apiKey {
|
||||||
|
http.Error(w, "missing auth", http.StatusUnauthorized)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
_, _ = w.Write([]byte(`{"data":[{"id":"custom-model"}]}`))
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
model := &config.ModelConfig{
|
||||||
|
Model: "openai/custom-model",
|
||||||
|
APIBase: srv.URL + "/v1",
|
||||||
|
}
|
||||||
|
model.SetAPIKey(apiKey)
|
||||||
|
|
||||||
|
if !probeLocalModelAvailability(model) {
|
||||||
|
t.Fatal("probeLocalModelAvailability() = false, want true when api_key is configured")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -36,11 +36,11 @@ func TestHandleListModels_ConfiguredStatusUsesRuntimeProbesForLocalModels(t *tes
|
||||||
var ollamaProbes []string
|
var ollamaProbes []string
|
||||||
var tcpProbes []string
|
var tcpProbes []string
|
||||||
|
|
||||||
probeOpenAICompatibleModelFunc = func(apiBase, modelID string) bool {
|
probeOpenAICompatibleModelFunc = func(apiBase, modelID, apiKey string) bool {
|
||||||
mu.Lock()
|
mu.Lock()
|
||||||
openAIProbes = append(openAIProbes, apiBase+"|"+modelID)
|
openAIProbes = append(openAIProbes, apiBase+"|"+modelID+"|"+apiKey)
|
||||||
mu.Unlock()
|
mu.Unlock()
|
||||||
return apiBase == "http://127.0.0.1:8000/v1" && modelID == "custom-model"
|
return apiBase == "http://127.0.0.1:8000/v1" && modelID == "custom-model" && apiKey == ""
|
||||||
}
|
}
|
||||||
probeOllamaModelFunc = func(apiBase, modelID string) bool {
|
probeOllamaModelFunc = func(apiBase, modelID string) bool {
|
||||||
mu.Lock()
|
mu.Lock()
|
||||||
|
|
@ -135,7 +135,7 @@ func TestHandleListModels_ConfiguredStatusUsesRuntimeProbesForLocalModels(t *tes
|
||||||
if !got["copilot-gpt-5.4"] {
|
if !got["copilot-gpt-5.4"] {
|
||||||
t.Fatalf("copilot model configured = false, want true when local bridge probe succeeds")
|
t.Fatalf("copilot model configured = false, want true when local bridge probe succeeds")
|
||||||
}
|
}
|
||||||
if len(openAIProbes) != 1 || openAIProbes[0] != "http://127.0.0.1:8000/v1|custom-model" {
|
if len(openAIProbes) != 1 || openAIProbes[0] != "http://127.0.0.1:8000/v1|custom-model|" {
|
||||||
t.Fatalf("openAI probes = %#v, want only local vllm probe", openAIProbes)
|
t.Fatalf("openAI probes = %#v, want only local vllm probe", openAIProbes)
|
||||||
}
|
}
|
||||||
if len(ollamaProbes) != 1 || ollamaProbes[0] != "http://localhost:11434/v1|llama3" {
|
if len(ollamaProbes) != 1 || ollamaProbes[0] != "http://localhost:11434/v1|llama3" {
|
||||||
|
|
@ -209,7 +209,7 @@ func TestHandleListModels_ProbesLocalModelsConcurrently(t *testing.T) {
|
||||||
started := make(chan string, 2)
|
started := make(chan string, 2)
|
||||||
release := make(chan struct{})
|
release := make(chan struct{})
|
||||||
|
|
||||||
probeOpenAICompatibleModelFunc = func(apiBase, modelID string) bool {
|
probeOpenAICompatibleModelFunc = func(apiBase, modelID, apiKey string) bool {
|
||||||
started <- apiBase + "|" + modelID
|
started <- apiBase + "|" + modelID
|
||||||
<-release
|
<-release
|
||||||
return true
|
return true
|
||||||
|
|
@ -269,9 +269,9 @@ func TestHandleListModels_NormalizesWildcardLocalAPIBaseForProbe(t *testing.T) {
|
||||||
resetModelProbeHooks(t)
|
resetModelProbeHooks(t)
|
||||||
|
|
||||||
var gotProbe string
|
var gotProbe string
|
||||||
probeOpenAICompatibleModelFunc = func(apiBase, modelID string) bool {
|
probeOpenAICompatibleModelFunc = func(apiBase, modelID, apiKey string) bool {
|
||||||
gotProbe = apiBase + "|" + modelID
|
gotProbe = apiBase + "|" + modelID + "|" + apiKey
|
||||||
return apiBase == "http://127.0.0.1:8000/v1" && modelID == "custom-model"
|
return apiBase == "http://127.0.0.1:8000/v1" && modelID == "custom-model" && apiKey == ""
|
||||||
}
|
}
|
||||||
|
|
||||||
cfg, err := config.LoadConfig(configPath)
|
cfg, err := config.LoadConfig(configPath)
|
||||||
|
|
@ -311,7 +311,7 @@ func TestHandleListModels_NormalizesWildcardLocalAPIBaseForProbe(t *testing.T) {
|
||||||
if !resp.Models[0].Configured {
|
if !resp.Models[0].Configured {
|
||||||
t.Fatal("wildcard-bound local model configured = false, want true after probe host normalization")
|
t.Fatal("wildcard-bound local model configured = false, want true after probe host normalization")
|
||||||
}
|
}
|
||||||
if gotProbe != "http://127.0.0.1:8000/v1|custom-model" {
|
if gotProbe != "http://127.0.0.1:8000/v1|custom-model|" {
|
||||||
t.Fatalf("probe api base = %q, want %q", gotProbe, "http://127.0.0.1:8000/v1|custom-model")
|
t.Fatalf("probe api base = %q, want %q", gotProbe, "http://127.0.0.1:8000/v1|custom-model|")
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue