feat(isolation): further hardening for agent loop and tools
This commit is contained in:
parent
999bf20144
commit
7e470c195a
4 changed files with 108 additions and 86 deletions
|
|
@ -195,8 +195,8 @@ func TestProcessMessage_IsolatedTenant_UsesPrivateWorkspace(t *testing.T) {
|
|||
}
|
||||
|
||||
// Verify history is in the base sessions directory with the isolated key
|
||||
// agent:::main:tenant-A becomes agent___main_tenant-A
|
||||
isoSessionPath := filepath.Join(tmpDir, "sessions", "agent___main_tenant-A.jsonl")
|
||||
// agent:main:tenant-A becomes agent_main_tenant-A
|
||||
isoSessionPath := filepath.Join(tmpDir, "sessions", "agent_main_tenant-A.jsonl")
|
||||
if _, err := os.Stat(isoSessionPath); os.IsNotExist(err) {
|
||||
t.Errorf("expected history at %s to exist", isoSessionPath)
|
||||
} else {
|
||||
|
|
|
|||
|
|
@ -113,7 +113,7 @@ const (
|
|||
toolLimitResponse = "I've reached `max_tool_iterations` without a final response. Increase `max_tool_iterations` in config.json if this task needs more tool steps."
|
||||
toolRepeatLoopResponse = "Detected repeated tool calls without progress; stopping to avoid an infinite loop."
|
||||
handledToolResponseSummary = "Requested output delivered via tool attachment."
|
||||
sessionKeyAgentPrefix = "agent::"
|
||||
sessionKeyAgentPrefix = "agent"
|
||||
metadataKeyAccountID = "account_id"
|
||||
metadataKeyGuildID = "guild_id"
|
||||
metadataKeyTeamID = "team_id"
|
||||
|
|
@ -1430,64 +1430,14 @@ func (al *AgentLoop) processMessage(ctx context.Context, msg bus.InboundMessage)
|
|||
return al.processSystemMessage(ctx, msg)
|
||||
}
|
||||
|
||||
route, baseAgent, routeErr := al.resolveMessageRoute(msg)
|
||||
route, _, routeErr := al.resolveMessageRoute(msg)
|
||||
if routeErr != nil {
|
||||
return "", routeErr
|
||||
}
|
||||
|
||||
agent := baseAgent
|
||||
isolationID := msg.ChatID
|
||||
if isolationID != "" && isolationID != "direct" {
|
||||
// Check agent instance cache first (keyed by channel:chatID)
|
||||
cacheKey := msg.Channel + ":" + isolationID
|
||||
if cached, ok := al.agentCache.Load(cacheKey); ok {
|
||||
agent = cached.(*AgentInstance)
|
||||
// Update last access time for TTL tracking
|
||||
al.lastCacheCheck.Store(cacheKey, time.Now())
|
||||
|
||||
logger.InfoCF("agent", "Reusing cached agent instance", map[string]any{
|
||||
"agent_id": agent.ID,
|
||||
"cache_key": cacheKey,
|
||||
"isolation_id": isolationID,
|
||||
})
|
||||
} else {
|
||||
// Create a transient isolated instance for this chat session
|
||||
// This ensures workspace, memory, and sessions are private to the chat_id.
|
||||
|
||||
// Determine the original config for this agent to preserve its specialized prompt/skills
|
||||
var ac *config.AgentConfig
|
||||
for i := range al.cfg.Agents.List {
|
||||
if routing.NormalizeAgentID(al.cfg.Agents.List[i].ID) == route.AgentID {
|
||||
ac = &al.cfg.Agents.List[i]
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
// Create a new instance with the isolationID
|
||||
// NewAgentInstance uses isolationID to sub-path the workspace
|
||||
agent = NewAgentInstance(ac, &al.cfg.Agents.Defaults, al.cfg, baseAgent.Provider, isolationID)
|
||||
|
||||
// Set its ID to match the routed agent so prompts and logs match
|
||||
agent.ID = route.AgentID
|
||||
|
||||
// Inject media store so tools (like send_file) can function
|
||||
agent.Tools.SetMediaStore(al.mediaStore)
|
||||
|
||||
// Re-register shared tools (web, message, spawn) to this transient agent
|
||||
// We pass a mini-registry containing only this agent
|
||||
registerSharedTools(al, al.cfg, al.bus, &AgentRegistry{agents: map[string]*AgentInstance{agent.ID: agent}}, baseAgent.Provider)
|
||||
|
||||
// Cache this agent instance per chat session
|
||||
al.agentCache.Store(cacheKey, agent)
|
||||
al.lastCacheCheck.Store(cacheKey, time.Now())
|
||||
|
||||
logger.InfoCF("agent", "Created isolated transient agent", map[string]any{
|
||||
"agent_id": agent.ID,
|
||||
"cache_key": cacheKey,
|
||||
"isolation_id": isolationID,
|
||||
"workspace": agent.Workspace,
|
||||
})
|
||||
}
|
||||
agent, err := al.getOrCreateIsolatedAgent(route.AgentID, msg.Channel, msg.ChatID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
// Reset message-tool state for this round so we don't skip publishing due to a previous round.
|
||||
|
|
@ -1609,6 +1559,68 @@ func (al *AgentLoop) requeueInboundMessage(msg bus.InboundMessage) error {
|
|||
})
|
||||
}
|
||||
|
||||
func (al *AgentLoop) getOrCreateIsolatedAgent(agentID, channel, isolationID string) (*AgentInstance, error) {
|
||||
if isolationID == "" || isolationID == "direct" {
|
||||
agent, ok := al.GetRegistry().GetAgent(agentID)
|
||||
if !ok {
|
||||
agent = al.GetRegistry().GetDefaultAgent()
|
||||
}
|
||||
if agent == nil {
|
||||
return nil, fmt.Errorf("no agent available for id %s", agentID)
|
||||
}
|
||||
return agent, nil
|
||||
}
|
||||
|
||||
cacheKey := channel + ":" + isolationID
|
||||
if cached, ok := al.agentCache.Load(cacheKey); ok {
|
||||
agent := cached.(*AgentInstance)
|
||||
al.lastCacheCheck.Store(cacheKey, time.Now())
|
||||
return agent, nil
|
||||
}
|
||||
|
||||
// Create a transient isolated instance for this chat session
|
||||
// This ensures workspace, memory, and sessions are private to the chat_id.
|
||||
|
||||
// Determine the original config for this agent to preserve its specialized prompt/skills
|
||||
var ac *config.AgentConfig
|
||||
for i := range al.cfg.Agents.List {
|
||||
if routing.NormalizeAgentID(al.cfg.Agents.List[i].ID) == agentID {
|
||||
ac = &al.cfg.Agents.List[i]
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
baseAgent, ok := al.GetRegistry().GetAgent(agentID)
|
||||
if !ok {
|
||||
baseAgent = al.GetRegistry().GetDefaultAgent()
|
||||
}
|
||||
if baseAgent == nil {
|
||||
return nil, fmt.Errorf("base agent %s not found", agentID)
|
||||
}
|
||||
|
||||
agent := NewAgentInstance(ac, &al.cfg.Agents.Defaults, al.cfg, baseAgent.Provider, isolationID)
|
||||
agent.ID = agentID
|
||||
|
||||
// Inject media store so tools (like send_file) can function
|
||||
agent.Tools.SetMediaStore(al.mediaStore)
|
||||
|
||||
// Re-register shared tools (web, message, spawn) to this transient agent
|
||||
registerSharedTools(al, al.cfg, al.bus, &AgentRegistry{agents: map[string]*AgentInstance{agent.ID: agent}}, baseAgent.Provider)
|
||||
|
||||
// Cache this agent instance per chat session
|
||||
al.agentCache.Store(cacheKey, agent)
|
||||
al.lastCacheCheck.Store(cacheKey, time.Now())
|
||||
|
||||
logger.InfoCF("agent", "Created isolated transient agent", map[string]any{
|
||||
"agent_id": agent.ID,
|
||||
"cache_key": cacheKey,
|
||||
"isolation_id": isolationID,
|
||||
"workspace": agent.Workspace,
|
||||
})
|
||||
|
||||
return agent, nil
|
||||
}
|
||||
|
||||
func (al *AgentLoop) processSystemMessage(
|
||||
ctx context.Context,
|
||||
msg bus.InboundMessage,
|
||||
|
|
@ -1654,14 +1666,18 @@ func (al *AgentLoop) processSystemMessage(
|
|||
return "", nil
|
||||
}
|
||||
|
||||
// Use default agent for system messages
|
||||
agent := al.GetRegistry().GetDefaultAgent()
|
||||
if agent == nil {
|
||||
return "", fmt.Errorf("no default agent for system message")
|
||||
// Use default agent for system messages, but lookup/create isolated tenant instances
|
||||
// that match the origin of the follow-up task. This ensures workspace isolation.
|
||||
agent, err := al.getOrCreateIsolatedAgent(routing.DefaultAgentID, originChannel, originChatID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
// Use the origin session for context
|
||||
sessionKey := routing.BuildAgentMainSessionKey(agent.ID)
|
||||
// Use provided session key if available, otherwise fall back to main
|
||||
sessionKey := msg.SessionKey
|
||||
if sessionKey == "" {
|
||||
sessionKey = routing.BuildAgentMainSessionKey(agent.ID)
|
||||
}
|
||||
|
||||
return al.runAgentLoop(ctx, agent, processOptions{
|
||||
SessionKey: sessionKey,
|
||||
|
|
@ -2357,7 +2373,8 @@ turnLoop:
|
|||
},
|
||||
)
|
||||
|
||||
llmResponseFields := map[string]any{
|
||||
logger.DebugCF("agent", "LLM response",
|
||||
map[string]any{
|
||||
"agent_id": ts.agent.ID,
|
||||
"iteration": iteration,
|
||||
"content_chars": len(response.Content),
|
||||
|
|
@ -2365,13 +2382,7 @@ turnLoop:
|
|||
"reasoning": response.Reasoning,
|
||||
"target_channel": al.targetReasoningChannelID(ts.channel),
|
||||
"channel": ts.channel,
|
||||
}
|
||||
if response.Usage != nil {
|
||||
llmResponseFields["prompt_tokens"] = response.Usage.PromptTokens
|
||||
llmResponseFields["completion_tokens"] = response.Usage.CompletionTokens
|
||||
llmResponseFields["total_tokens"] = response.Usage.TotalTokens
|
||||
}
|
||||
logger.DebugCF("agent", "LLM response", llmResponseFields)
|
||||
})
|
||||
|
||||
if len(response.ToolCalls) == 0 || gracefulTerminal {
|
||||
responseContent := response.Content
|
||||
|
|
@ -2668,6 +2679,7 @@ turnLoop:
|
|||
SenderID: fmt.Sprintf("async:%s", asyncToolName),
|
||||
ChatID: fmt.Sprintf("%s:%s", ts.channel, ts.chatID),
|
||||
Content: content,
|
||||
SessionKey: ts.opts.SessionKey,
|
||||
})
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1400,7 +1400,7 @@ func TestProcessMessage_UsesRouteSessionKey(t *testing.T) {
|
|||
}
|
||||
|
||||
// With chatID isolation, session key is derived from chatID
|
||||
sessionKey := fmt.Sprintf("agent:::main:%s", msg.ChatID)
|
||||
sessionKey := fmt.Sprintf("agent:main:%s", msg.ChatID)
|
||||
|
||||
defaultAgent := al.registry.GetDefaultAgent()
|
||||
if defaultAgent == nil {
|
||||
|
|
|
|||
|
|
@ -1061,18 +1061,28 @@ func (t *ExecTool) guardCommand(command, cwd string) string {
|
|||
// Web URL schemes whose path components (starting with //) should be exempt
|
||||
// from workspace sandbox checks. file: is intentionally excluded so that
|
||||
// file:// URIs are still validated against the workspace boundary.
|
||||
webSchemes := []string{"http:", "https:", "ftp:", "ftps:", "sftp:", "ssh:", "git:"}
|
||||
webSchemes := []string{"http:", "https:", "ftp:", "ftps:", "ssh:", "git:", "sftp:"}
|
||||
|
||||
matchIndices := absolutePathPattern.FindAllStringIndex(cmd, -1)
|
||||
|
||||
for _, loc := range matchIndices {
|
||||
raw := cmd[loc[0]:loc[1]]
|
||||
|
||||
// Check if this is truly the start of a path component.
|
||||
// It should be at the start of the command or preceded by a shell delimiter.
|
||||
if loc[0] > 0 {
|
||||
prev := cmd[loc[0]-1]
|
||||
// Typical shell delimiters that separate command arguments or environment variables.
|
||||
// We include space-like chars, basic separators, and assignment equals.
|
||||
// We also include ':' because it precedes paths in lists ($PATH) and URLs (file://, https://).
|
||||
if !strings.ContainsAny(string(prev), " \t\n\r;|\"&!<>(){}=[]':") {
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
// Skip URL path components that look like they're from web URLs.
|
||||
// When a URL like "https://github.com" is parsed, the regex captures
|
||||
// "//github.com" as a match (the path portion after "https:").
|
||||
// Use the exact match position (loc[0]) so that duplicate //path substrings
|
||||
// in the same command are each evaluated at their own position.
|
||||
if strings.HasPrefix(raw, "//") && loc[0] > 0 {
|
||||
before := cmd[:loc[0]]
|
||||
isWebURL := false
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue