feat: add workspace path validation for sending attachments

This commit is contained in:
XZB-1248 2026-02-19 22:18:26 +08:00
parent a31403c547
commit 6d4d2ed7a0
3 changed files with 58 additions and 40 deletions

View file

@ -15,20 +15,21 @@ import (
// AgentInstance represents a fully configured agent with its own workspace, // AgentInstance represents a fully configured agent with its own workspace,
// session manager, context builder, and tool registry. // session manager, context builder, and tool registry.
type AgentInstance struct { type AgentInstance struct {
ID string ID string
Name string Name string
Model string Model string
Fallbacks []string Fallbacks []string
Workspace string Workspace string
MaxIterations int MaxIterations int
ContextWindow int ContextWindow int
Provider providers.LLMProvider RestrictToWorkspace bool
Sessions *session.SessionManager Provider providers.LLMProvider
ContextBuilder *ContextBuilder Sessions *session.SessionManager
Tools *tools.ToolRegistry ContextBuilder *ContextBuilder
Subagents *config.SubagentsConfig Tools *tools.ToolRegistry
SkillsFilter []string Subagents *config.SubagentsConfig
Candidates []providers.FallbackCandidate SkillsFilter []string
Candidates []providers.FallbackCandidate
} }
// NewAgentInstance creates an agent instance from config. // NewAgentInstance creates an agent instance from config.
@ -84,20 +85,21 @@ func NewAgentInstance(
candidates := providers.ResolveCandidates(modelCfg, defaults.Provider) candidates := providers.ResolveCandidates(modelCfg, defaults.Provider)
return &AgentInstance{ return &AgentInstance{
ID: agentID, ID: agentID,
Name: agentName, Name: agentName,
Model: model, Model: model,
Fallbacks: fallbacks, Fallbacks: fallbacks,
Workspace: workspace, Workspace: workspace,
MaxIterations: maxIter, MaxIterations: maxIter,
ContextWindow: defaults.MaxTokens, ContextWindow: defaults.MaxTokens,
Provider: provider, RestrictToWorkspace: restrict,
Sessions: sessionsManager, Provider: provider,
ContextBuilder: contextBuilder, Sessions: sessionsManager,
Tools: toolsRegistry, ContextBuilder: contextBuilder,
Subagents: subagents, Tools: toolsRegistry,
SkillsFilter: skillsFilter, Subagents: subagents,
Candidates: candidates, SkillsFilter: skillsFilter,
Candidates: candidates,
} }
} }

View file

@ -106,7 +106,7 @@ func registerSharedTools(cfg *config.Config, msgBus *bus.MessageBus, registry *A
agent.Tools.Register(tools.NewSPITool()) agent.Tools.Register(tools.NewSPITool())
// Message tool // Message tool
messageTool := tools.NewMessageTool() messageTool := tools.NewMessageTool(agent.Workspace, agent.RestrictToWorkspace)
messageTool.SetSendCallback(func(channel, chatID, content string, attachments []bus.Attachment) error { messageTool.SetSendCallback(func(channel, chatID, content string, attachments []bus.Attachment) error {
msgBus.PublishOutbound(bus.OutboundMessage{ msgBus.PublishOutbound(bus.OutboundMessage{
Channel: channel, Channel: channel,

View file

@ -10,14 +10,20 @@ import (
type SendCallback func(channel, chatID, content string, attachments []bus.Attachment) error type SendCallback func(channel, chatID, content string, attachments []bus.Attachment) error
type MessageTool struct { type MessageTool struct {
allowedDir string
restrict bool
sendCallback SendCallback sendCallback SendCallback
defaultChannel string defaultChannel string
defaultChatID string defaultChatID string
sentInRound bool // Tracks whether a message was sent in the current processing round sentInRound bool // Tracks whether a message was sent in the current processing round
} }
func NewMessageTool() *MessageTool { // NewMessageTool creates a new MessageTool with optional uploading directory restriction.
return &MessageTool{} func NewMessageTool(allowedDir string, restrict bool) *MessageTool {
return &MessageTool{
allowedDir: allowedDir,
restrict: restrict,
}
} }
func (t *MessageTool) Name() string { func (t *MessageTool) Name() string {
@ -110,16 +116,26 @@ func (t *MessageTool) Execute(ctx context.Context, args map[string]interface{})
var attachments []bus.Attachment var attachments []bus.Attachment
if attachmentsRaw, ok := args["attachments"].([]interface{}); ok { if attachmentsRaw, ok := args["attachments"].([]interface{}); ok {
for _, attachRaw := range attachmentsRaw { for _, attachRaw := range attachmentsRaw {
if attachMap, ok := attachRaw.(map[string]interface{}); ok { attachMap, ok := attachRaw.(map[string]interface{})
path, pathOk := attachMap["path"].(string) if !ok {
filename, filenameOk := attachMap["filename"].(string) continue // Skip invalid attachment entries
if pathOk && filenameOk {
attachments = append(attachments, bus.Attachment{
Path: path,
Filename: filename,
})
}
} }
path, pathOk := attachMap["path"].(string)
filename, filenameOk := attachMap["filename"].(string)
if !pathOk || !filenameOk {
continue // Skip invalid attachment entries
}
resolvedPath, err := validatePath(path, t.allowedDir, t.restrict)
if err != nil {
return ErrorResult(err.Error())
}
attachments = append(attachments, bus.Attachment{
Path: resolvedPath,
Filename: filename,
})
} }
} }