feat: add workspace path validation for sending attachments
This commit is contained in:
parent
a31403c547
commit
6d4d2ed7a0
3 changed files with 58 additions and 40 deletions
|
|
@ -22,6 +22,7 @@ type AgentInstance struct {
|
|||
Workspace string
|
||||
MaxIterations int
|
||||
ContextWindow int
|
||||
RestrictToWorkspace bool
|
||||
Provider providers.LLMProvider
|
||||
Sessions *session.SessionManager
|
||||
ContextBuilder *ContextBuilder
|
||||
|
|
@ -91,6 +92,7 @@ func NewAgentInstance(
|
|||
Workspace: workspace,
|
||||
MaxIterations: maxIter,
|
||||
ContextWindow: defaults.MaxTokens,
|
||||
RestrictToWorkspace: restrict,
|
||||
Provider: provider,
|
||||
Sessions: sessionsManager,
|
||||
ContextBuilder: contextBuilder,
|
||||
|
|
|
|||
|
|
@ -106,7 +106,7 @@ func registerSharedTools(cfg *config.Config, msgBus *bus.MessageBus, registry *A
|
|||
agent.Tools.Register(tools.NewSPITool())
|
||||
|
||||
// Message tool
|
||||
messageTool := tools.NewMessageTool()
|
||||
messageTool := tools.NewMessageTool(agent.Workspace, agent.RestrictToWorkspace)
|
||||
messageTool.SetSendCallback(func(channel, chatID, content string, attachments []bus.Attachment) error {
|
||||
msgBus.PublishOutbound(bus.OutboundMessage{
|
||||
Channel: channel,
|
||||
|
|
|
|||
|
|
@ -10,14 +10,20 @@ import (
|
|||
type SendCallback func(channel, chatID, content string, attachments []bus.Attachment) error
|
||||
|
||||
type MessageTool struct {
|
||||
allowedDir string
|
||||
restrict bool
|
||||
sendCallback SendCallback
|
||||
defaultChannel string
|
||||
defaultChatID string
|
||||
sentInRound bool // Tracks whether a message was sent in the current processing round
|
||||
}
|
||||
|
||||
func NewMessageTool() *MessageTool {
|
||||
return &MessageTool{}
|
||||
// NewMessageTool creates a new MessageTool with optional uploading directory restriction.
|
||||
func NewMessageTool(allowedDir string, restrict bool) *MessageTool {
|
||||
return &MessageTool{
|
||||
allowedDir: allowedDir,
|
||||
restrict: restrict,
|
||||
}
|
||||
}
|
||||
|
||||
func (t *MessageTool) Name() string {
|
||||
|
|
@ -110,18 +116,28 @@ func (t *MessageTool) Execute(ctx context.Context, args map[string]interface{})
|
|||
var attachments []bus.Attachment
|
||||
if attachmentsRaw, ok := args["attachments"].([]interface{}); ok {
|
||||
for _, attachRaw := range attachmentsRaw {
|
||||
if attachMap, ok := attachRaw.(map[string]interface{}); ok {
|
||||
attachMap, ok := attachRaw.(map[string]interface{})
|
||||
if !ok {
|
||||
continue // Skip invalid attachment entries
|
||||
}
|
||||
|
||||
path, pathOk := attachMap["path"].(string)
|
||||
filename, filenameOk := attachMap["filename"].(string)
|
||||
if pathOk && filenameOk {
|
||||
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: path,
|
||||
Path: resolvedPath,
|
||||
Filename: filename,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if err := t.sendCallback(channel, chatID, content, attachments); err != nil {
|
||||
return &ToolResult{
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue