Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 11 additions & 2 deletions agent/agent.go
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,9 @@ type Options struct {
// tool is still registered but returns "[no human available]" instead of
// blocking (e.g. standalone squadron with no commander attached).
HumanBridge aitools.HumanInputBridge
// GatewayBridge powers the `builtins.gateway.post` tool. When nil, the
// tool is still registered but returns "[no gateway configured]".
GatewayBridge aitools.GatewayBridge
}

// New creates a new agent from config
Expand Down Expand Up @@ -155,7 +158,7 @@ func New(ctx context.Context, opts Options) (*Agent, error) {

// Build tools map and add sanitized aliases so LLM tool calls
// (which use API-safe names like "plugins_shell_echo") resolve correctly
tools := config.BuildToolsMap(agentCfg.Tools, cfg.CustomTools, cfg.LoadedPlugins, cfg.LoadedMCPClients, opts.DatasetStore, opts.HumanBridge)
tools := config.BuildToolsMap(agentCfg.Tools, cfg.CustomTools, cfg.LoadedPlugins, cfg.LoadedMCPClients, opts.DatasetStore, opts.HumanBridge, opts.GatewayBridge)
aitools.AddSanitizedAliases(tools)

// Create result store and interceptor for large results
Expand Down Expand Up @@ -186,6 +189,12 @@ func New(ctx context.Context, opts Options) (*Agent, error) {
tools["file_delete"] = &aitools.MemoryDeleteTool{Store: opts.MemoryStore}
tools["file_search"] = &aitools.MemorySearchTool{Store: opts.MemoryStore}
tools["file_grep"] = &aitools.MemoryGrepTool{Store: opts.MemoryStore}
// The gateway post tool resolves attachments from the same store.
for _, tool := range tools {
if gp, ok := tool.(*aitools.GatewayPostTool); ok {
gp.Store = opts.MemoryStore
}
}
}

// Resolve skills and add load_skill tool
Expand All @@ -197,7 +206,7 @@ func New(ctx context.Context, opts Options) (*Agent, error) {
AvailableSkills: availableSkills,
AgentTools: tools,
ToolBuilder: func(toolRefs []string) map[string]aitools.Tool {
t := config.BuildToolsMap(toolRefs, cfg.CustomTools, cfg.LoadedPlugins, cfg.LoadedMCPClients, opts.DatasetStore, opts.HumanBridge)
t := config.BuildToolsMap(toolRefs, cfg.CustomTools, cfg.LoadedPlugins, cfg.LoadedMCPClients, opts.DatasetStore, opts.HumanBridge, opts.GatewayBridge)
aitools.AddSanitizedAliases(t)
return t
},
Expand Down
5 changes: 5 additions & 0 deletions agent/agent_manager.go
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ type AgentManager struct {
provider llm.Provider // optional injected provider for agents
budget BudgetChecker
humanBridge aitools.HumanInputBridge // bridge for builtins.human.ask on spawned agents
gatewayBridge aitools.GatewayBridge // bridge for builtins.gateway.post on spawned agents
}

// AgentManagerConfig holds the dependencies needed to create an AgentManager.
Expand All @@ -70,6 +71,8 @@ type AgentManagerConfig struct {
Budget BudgetChecker
// HumanBridge — nil disables builtins.human.ask on spawned agents.
HumanBridge aitools.HumanInputBridge
// GatewayBridge — nil disables builtins.gateway.post on spawned agents.
GatewayBridge aitools.GatewayBridge
}

// NewAgentManager creates a new AgentManager.
Expand All @@ -95,6 +98,7 @@ func NewAgentManager(cfg AgentManagerConfig) *AgentManager {
provider: cfg.Provider,
budget: cfg.Budget,
humanBridge: cfg.HumanBridge,
gatewayBridge: cfg.GatewayBridge,
}
}

Expand Down Expand Up @@ -284,6 +288,7 @@ func (m *AgentManager) createAgent(ctx context.Context, agentCfg *config.Agent)
PricingOverrides: m.pricingOverrides,
Budget: m.budget,
HumanBridge: m.humanBridge,
GatewayBridge: m.gatewayBridge,
})
}

Expand Down
11 changes: 11 additions & 0 deletions agent/commander.go
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,9 @@ type CommanderOptions struct {
// spawns. Nil disables HITL — the tool then returns
// "[no human available]" instead of blocking.
HumanBridge aitools.HumanInputBridge
// GatewayBridge powers builtins.gateway.post on agents this commander
// spawns. Nil → the tool returns "[no gateway configured]".
GatewayBridge aitools.GatewayBridge
}

// DependencyOutputSchema describes a completed dependency task's output schema
Expand Down Expand Up @@ -348,6 +351,7 @@ type Commander struct {
pruneTo int // Prune down to this many turns
budget BudgetChecker // Optional token/dollar budget enforcer
humanBridge aitools.HumanInputBridge // Optional bridge for builtins.human.ask
gatewayBridge aitools.GatewayBridge // Optional bridge for builtins.gateway.post
}

// NewCommander creates a new commander for a mission task
Expand Down Expand Up @@ -473,6 +477,7 @@ func NewCommander(ctx context.Context, opts CommanderOptions) (*Commander, error
pricingOverrides: opts.PricingOverrides,
budget: opts.Budget,
humanBridge: opts.HumanBridge,
gatewayBridge: opts.GatewayBridge,
}

// Add result tools to commander's tool map
Expand All @@ -491,6 +496,11 @@ func NewCommander(ctx context.Context, opts CommanderOptions) (*Commander, error
sup.tools["file_delete"] = &aitools.MemoryDeleteTool{Store: opts.MemoryStore}
sup.tools["file_search"] = &aitools.MemorySearchTool{Store: opts.MemoryStore}
sup.tools["file_grep"] = &aitools.MemoryGrepTool{Store: opts.MemoryStore}
for _, tool := range sup.tools {
if gp, ok := tool.(*aitools.GatewayPostTool); ok {
gp.Store = opts.MemoryStore
}
}
if memoryPrompt := prompts.FormatMemoryContext(opts.MemoryStore); memoryPrompt != "" {
session.AddSystemPrompt(memoryPrompt)
}
Expand Down Expand Up @@ -737,6 +747,7 @@ func (s *Commander) SetToolCallbacks(callbacks *CommanderToolCallbacks, depSumma
Provider: s.provider,
Budget: s.budget,
HumanBridge: s.humanBridge,
GatewayBridge: s.gatewayBridge,
})
}

Expand Down
248 changes: 248 additions & 0 deletions aitools/gateway_post.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,248 @@
package aitools

import (
"context"
"encoding/json"
"fmt"
"mime"
"net/http"
"os"
"path/filepath"
"strings"
)

// GatewayBridge posts a message through the configured gateway subprocess
// (Discord, Slack, …) and advertises how that gateway wants messages shaped.
// Pass nil to build a tool that reports the gateway is unavailable instead of
// posting — the tool is always registered.
type GatewayBridge interface {
// PostMessage forwards the raw, gateway-schema-shaped JSON the agent
// produced (text + rich layout) plus any squadron-resolved file
// attachments. The gateway parses the payload and uploads the attachment
// bytes directly.
PostMessage(ctx context.Context, payload string, attachments []GatewayAttachment) error
// MessageToolDescription is the gateway-supplied tool description (how to
// format messages for this gateway). Empty → squadron's default.
MessageToolDescription() string
// MessageToolSchema is the gateway-supplied JSON Schema for the tool's
// params. Empty → squadron's default { message, channel } shape.
MessageToolSchema() string
}

// GatewayAttachment is a squadron-local file resolved from the mission's
// memory/scratchpad/packet storage, shipped to the gateway as raw bytes.
type GatewayAttachment struct {
Filename string
MimeType string
Content []byte
}

// Attachments are sourced from squadron-local files only (never a URL the
// model picks), so there is no SSRF surface. Caps keep a single post within
// the gateway gRPC channel's message-size budget.
const (
maxAttachmentBytes = 25 << 20 // 25 MiB per file
maxTotalAttachmentBytes = 30 << 20 // 30 MiB per post
)

// GatewayPostTool backs builtins.gateway.post. The gateway owns the message
// contract (text + rich layout): its description and JSON Schema are surfaced
// to the LLM and the params are forwarded verbatim. Squadron owns the
// `attachments` field — it resolves each {slot, path} reference against the
// mission's MemoryStore and ships the bytes to the gateway.
type GatewayPostTool struct {
Bridge GatewayBridge
Store MemoryStore
}

func (t *GatewayPostTool) ToolName() string { return "post" }

const defaultGatewayPostDescription = "Post a message to the configured gateway's external system (Discord, Slack, etc.). " +
"If no gateway is configured, returns \"[no gateway configured]\" so you can proceed without failing."

const attachmentsDescriptionSuffix = "To attach files, set `attachments` to a list of {\"slot\":..., \"path\":...} objects " +
"referencing squadron's own memory/scratchpad/packet storage (NOT URLs) — squadron reads each file and uploads it."

const defaultGatewayPostSchema = `{
"type": "object",
"properties": {
"message": {"type": "string", "description": "The message text to post."},
"channel": {"type": "string", "description": "Optional channel name or id override."}
},
"required": ["message"]
}`

const attachmentsSchemaProperty = `{
"type": "array",
"description": "Optional files to attach, sourced from squadron's own memory/scratchpad/packet storage (NOT URLs). Each item references a local file by slot and path.",
"items": {
"type": "object",
"properties": {
"slot": {"type": "string", "description": "Slot: \"memory\", \"scratchpad\", a shared-memory name, or \"packet.<name>\"."},
"path": {"type": "string", "description": "Relative path within the slot, e.g. \"report.pdf\"."}
},
"required": ["slot", "path"]
}
}`

type gatewayAttachmentRef struct {
Slot string `json:"slot"`
Path string `json:"path"`
}

func (t *GatewayPostTool) ToolDescription() string {
base := defaultGatewayPostDescription
if t.Bridge != nil {
if d := t.Bridge.MessageToolDescription(); d != "" {
base = d
}
}
if t.Store != nil {
return base + " " + attachmentsDescriptionSuffix
}
return base
}

func (t *GatewayPostTool) ToolPayloadSchema() Schema {
raw := ""
if t.Bridge != nil {
raw = t.Bridge.MessageToolSchema()
}
if strings.TrimSpace(raw) == "" {
raw = defaultGatewayPostSchema
}
// The gateway owns text + rich layout; squadron owns attachments (local
// files), so inject that field only when a memory store is available.
if t.Store != nil {
raw = injectAttachmentsProperty(raw)
}
return Schema{Type: TypeObject, Properties: PropertyMap{}}.WithRawJSONSchema(json.RawMessage(raw))
}

// injectAttachmentsProperty adds squadron's `attachments` property to the
// gateway-owned schema. Best-effort: if the schema can't be parsed or already
// defines `attachments`, it is returned unchanged.
func injectAttachmentsProperty(schema string) string {
var root map[string]json.RawMessage
if err := json.Unmarshal([]byte(schema), &root); err != nil {
return schema
}
props := map[string]json.RawMessage{}
if rawProps, ok := root["properties"]; ok {
if err := json.Unmarshal(rawProps, &props); err != nil {
return schema
}
}
if _, exists := props["attachments"]; exists {
return schema
}
props["attachments"] = json.RawMessage(attachmentsSchemaProperty)
newProps, err := json.Marshal(props)
if err != nil {
return schema
}
root["properties"] = newProps
out, err := json.Marshal(root)
if err != nil {
return schema
}
return string(out)
}

// NoGatewayObservation is what Call returns when no gateway bridge is wired.
const NoGatewayObservation = "[no gateway configured]"

func (t *GatewayPostTool) Call(ctx context.Context, params string) string {
if t.Bridge == nil {
return NoGatewayObservation
}
if s := strings.TrimSpace(params); s == "" || s == "{}" {
return "Error: empty message payload"
}
if !json.Valid([]byte(params)) {
return "Error: invalid JSON parameters"
}

payload := params
var attachments []GatewayAttachment

// Pull `attachments` out of the payload and resolve it against the mission's
// local file storage; the gateway never sees the references, only bytes.
var root map[string]json.RawMessage
if err := json.Unmarshal([]byte(params), &root); err != nil {
return "Error: invalid JSON parameters"
}
if rawAtt, ok := root["attachments"]; ok {
delete(root, "attachments")
var refs []gatewayAttachmentRef
if err := json.Unmarshal(rawAtt, &refs); err != nil {
return "Error: attachments must be an array of {slot, path} objects"
}
resolved, errMsg := t.resolveAttachments(refs)
if errMsg != "" {
return "Error: " + errMsg
}
attachments = resolved
rest, err := json.Marshal(root)
if err != nil {
return "Error: " + err.Error()
}
payload = string(rest)
}

if err := t.Bridge.PostMessage(ctx, payload, attachments); err != nil {
return "Error: " + err.Error()
}
return "Message posted to the gateway."
}

func (t *GatewayPostTool) resolveAttachments(refs []gatewayAttachmentRef) ([]GatewayAttachment, string) {
if len(refs) == 0 {
return nil, ""
}
if t.Store == nil {
return nil, "attachments are not available for this mission (no memory or scratchpad configured)"
}
var out []GatewayAttachment
var total int
for _, r := range refs {
if strings.TrimSpace(r.Slot) == "" || strings.TrimSpace(r.Path) == "" {
return nil, "each attachment needs a non-empty slot and path"
}
abs, err := resolveSlotPath(t.Store, r.Slot, r.Path)
if err != nil {
return nil, fmt.Sprintf("attachment %s/%s: %v", r.Slot, r.Path, err)
}
info, err := os.Stat(abs)
if err != nil {
return nil, fmt.Sprintf("attachment %s/%s: %v", r.Slot, r.Path, err)
}
if info.IsDir() {
return nil, fmt.Sprintf("attachment %s/%s: is a directory, not a file", r.Slot, r.Path)
}
if info.Size() > maxAttachmentBytes {
return nil, fmt.Sprintf("attachment %s/%s: %d bytes exceeds the %d-byte per-file limit", r.Slot, r.Path, info.Size(), maxAttachmentBytes)
}
data, err := os.ReadFile(abs)
if err != nil {
return nil, fmt.Sprintf("attachment %s/%s: %v", r.Slot, r.Path, err)
}
total += len(data)
if total > maxTotalAttachmentBytes {
return nil, fmt.Sprintf("total attachment size exceeds the %d-byte per-post limit", maxTotalAttachmentBytes)
}
out = append(out, GatewayAttachment{
Filename: filepath.Base(r.Path),
MimeType: detectAttachmentMime(r.Path, data),
Content: data,
})
}
return out, ""
}

func detectAttachmentMime(name string, data []byte) string {
if ct := mime.TypeByExtension(filepath.Ext(name)); ct != "" {
return ct
}
return http.DetectContentType(data)
}
Loading