feat(isolation): further hardening for agent loop and tools

This commit is contained in:
stevef 2026-03-27 22:50:21 +01:00
parent f2ffc6cc31
commit 6c277610df
4 changed files with 108 additions and 86 deletions

View file

@ -195,8 +195,8 @@ func TestProcessMessage_IsolatedTenant_UsesPrivateWorkspace(t *testing.T) {
} }
// Verify history is in the base sessions directory with the isolated key // Verify history is in the base sessions directory with the isolated key
// agent:::main:tenant-A becomes agent___main_tenant-A // agent:main:tenant-A becomes agent_main_tenant-A
isoSessionPath := filepath.Join(tmpDir, "sessions", "agent___main_tenant-A.jsonl") isoSessionPath := filepath.Join(tmpDir, "sessions", "agent_main_tenant-A.jsonl")
if _, err := os.Stat(isoSessionPath); os.IsNotExist(err) { if _, err := os.Stat(isoSessionPath); os.IsNotExist(err) {
t.Errorf("expected history at %s to exist", isoSessionPath) t.Errorf("expected history at %s to exist", isoSessionPath)
} else { } else {

View file

@ -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." 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." toolRepeatLoopResponse = "Detected repeated tool calls without progress; stopping to avoid an infinite loop."
handledToolResponseSummary = "Requested output delivered via tool attachment." handledToolResponseSummary = "Requested output delivered via tool attachment."
sessionKeyAgentPrefix = "agent::" sessionKeyAgentPrefix = "agent"
metadataKeyAccountID = "account_id" metadataKeyAccountID = "account_id"
metadataKeyGuildID = "guild_id" metadataKeyGuildID = "guild_id"
metadataKeyTeamID = "team_id" metadataKeyTeamID = "team_id"
@ -1430,64 +1430,14 @@ func (al *AgentLoop) processMessage(ctx context.Context, msg bus.InboundMessage)
return al.processSystemMessage(ctx, msg) return al.processSystemMessage(ctx, msg)
} }
route, baseAgent, routeErr := al.resolveMessageRoute(msg) route, _, routeErr := al.resolveMessageRoute(msg)
if routeErr != nil { if routeErr != nil {
return "", routeErr return "", routeErr
} }
agent := baseAgent agent, err := al.getOrCreateIsolatedAgent(route.AgentID, msg.Channel, msg.ChatID)
isolationID := msg.ChatID if err != nil {
if isolationID != "" && isolationID != "direct" { return "", err
// 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,
})
}
} }
// Reset message-tool state for this round so we don't skip publishing due to a previous round. // 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( func (al *AgentLoop) processSystemMessage(
ctx context.Context, ctx context.Context,
msg bus.InboundMessage, msg bus.InboundMessage,
@ -1654,14 +1666,18 @@ func (al *AgentLoop) processSystemMessage(
return "", nil return "", nil
} }
// Use default agent for system messages // Use default agent for system messages, but lookup/create isolated tenant instances
agent := al.GetRegistry().GetDefaultAgent() // that match the origin of the follow-up task. This ensures workspace isolation.
if agent == nil { agent, err := al.getOrCreateIsolatedAgent(routing.DefaultAgentID, originChannel, originChatID)
return "", fmt.Errorf("no default agent for system message") if err != nil {
return "", err
} }
// Use the origin session for context // Use provided session key if available, otherwise fall back to main
sessionKey := routing.BuildAgentMainSessionKey(agent.ID) sessionKey := msg.SessionKey
if sessionKey == "" {
sessionKey = routing.BuildAgentMainSessionKey(agent.ID)
}
return al.runAgentLoop(ctx, agent, processOptions{ return al.runAgentLoop(ctx, agent, processOptions{
SessionKey: sessionKey, SessionKey: sessionKey,
@ -2357,21 +2373,16 @@ turnLoop:
}, },
) )
llmResponseFields := map[string]any{ logger.DebugCF("agent", "LLM response",
"agent_id": ts.agent.ID, map[string]any{
"iteration": iteration, "agent_id": ts.agent.ID,
"content_chars": len(response.Content), "iteration": iteration,
"tool_calls": len(response.ToolCalls), "content_chars": len(response.Content),
"reasoning": response.Reasoning, "tool_calls": len(response.ToolCalls),
"target_channel": al.targetReasoningChannelID(ts.channel), "reasoning": response.Reasoning,
"channel": ts.channel, "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 { if len(response.ToolCalls) == 0 || gracefulTerminal {
responseContent := response.Content responseContent := response.Content
@ -2664,10 +2675,11 @@ turnLoop:
pubCtx, pubCancel := context.WithTimeout(context.Background(), 5*time.Second) pubCtx, pubCancel := context.WithTimeout(context.Background(), 5*time.Second)
defer pubCancel() defer pubCancel()
_ = al.bus.PublishInbound(pubCtx, bus.InboundMessage{ _ = al.bus.PublishInbound(pubCtx, bus.InboundMessage{
Channel: "system", Channel: "system",
SenderID: fmt.Sprintf("async:%s", asyncToolName), SenderID: fmt.Sprintf("async:%s", asyncToolName),
ChatID: fmt.Sprintf("%s:%s", ts.channel, ts.chatID), ChatID: fmt.Sprintf("%s:%s", ts.channel, ts.chatID),
Content: content, Content: content,
SessionKey: ts.opts.SessionKey,
}) })
} }

View file

@ -1400,7 +1400,7 @@ func TestProcessMessage_UsesRouteSessionKey(t *testing.T) {
} }
// With chatID isolation, session key is derived from chatID // 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() defaultAgent := al.registry.GetDefaultAgent()
if defaultAgent == nil { if defaultAgent == nil {

View file

@ -1061,18 +1061,28 @@ func (t *ExecTool) guardCommand(command, cwd string) string {
// Web URL schemes whose path components (starting with //) should be exempt // Web URL schemes whose path components (starting with //) should be exempt
// from workspace sandbox checks. file: is intentionally excluded so that // from workspace sandbox checks. file: is intentionally excluded so that
// file:// URIs are still validated against the workspace boundary. // 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) matchIndices := absolutePathPattern.FindAllStringIndex(cmd, -1)
for _, loc := range matchIndices { for _, loc := range matchIndices {
raw := cmd[loc[0]:loc[1]] 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. // Skip URL path components that look like they're from web URLs.
// When a URL like "https://github.com" is parsed, the regex captures // When a URL like "https://github.com" is parsed, the regex captures
// "//github.com" as a match (the path portion after "https:"). // "//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 { if strings.HasPrefix(raw, "//") && loc[0] > 0 {
before := cmd[:loc[0]] before := cmd[:loc[0]]
isWebURL := false isWebURL := false