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
// 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 {

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."
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,
})
}

View file

@ -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 {

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
// 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