From 33f975edf0bb7940bb71c163af3ab0ae7652f44c Mon Sep 17 00:00:00 2001
From: viettranx <viettranx@gmail.com>
Date: Mon, 30 Mar 2026 22:28:00 +0700
Subject: [PATCH] fix: use agent_key instead of UUID for session tool
 authorization (#573)

Session tools (sessions_list, session_status, sessions_history, sessions_send)
were using resolveAgentIDString(ctx) which returns the agent UUID, but session
keys are built using agent_key. This caused all session tool operations to fail
silently or return "access denied" for every channel.

Replace resolveAgentIDString() with ToolAgentKeyFromCtx() in all four session
tools. Add fail-closed guard for empty agent key.
---
 internal/tools/sessions.go         | 16 +++++++++-------
 internal/tools/sessions_history.go |  9 +++++----
 internal/tools/sessions_send.go    | 11 ++++++-----
 internal/tools/sessions_test.go    |  5 ++++-
 4 files changed, 24 insertions(+), 17 deletions(-)

diff --git a/internal/tools/sessions.go b/internal/tools/sessions.go
index 124198a6..8e8736e7 100644
--- a/internal/tools/sessions.go
+++ b/internal/tools/sessions.go
@@ -58,11 +58,12 @@ func (t *SessionsListTool) Execute(ctx context.Context, args map[string]any) *Re
 		activeMinutes = int(v)
 	}
 
-	agentID := resolveAgentIDString(ctx)
-	if agentID == "" {
+	// Session keys use agent_key (e.g. "agent:victoria:..."), not UUID.
+	agentKey := ToolAgentKeyFromCtx(ctx)
+	if agentKey == "" {
 		return ErrorResult("agent context required")
 	}
-	sessions := t.sessions.List(ctx, agentID)
+	sessions := t.sessions.List(ctx, agentKey)
 
 	// Filter by active_minutes
 	if activeMinutes > 0 {
@@ -145,12 +146,13 @@ func (t *SessionStatusTool) Execute(ctx context.Context, args map[string]any) *R
 		return ErrorResult("session_key is required (could not detect current session)")
 	}
 
-	// Security: validate session belongs to current agent (fail-closed)
-	agentID := resolveAgentIDString(ctx)
-	if agentID == "" {
+	// Security: validate session belongs to current agent (fail-closed).
+	// Session keys use agent_key (e.g. "agent:victoria:..."), not UUID.
+	agentKey := ToolAgentKeyFromCtx(ctx)
+	if agentKey == "" {
 		return ErrorResult("agent context required")
 	}
-	if !strings.HasPrefix(sessionKey, "agent:"+agentID+":") {
+	if !strings.HasPrefix(sessionKey, "agent:"+agentKey+":") {
 		return ErrorResult("access denied: session belongs to a different agent")
 	}
 
diff --git a/internal/tools/sessions_history.go b/internal/tools/sessions_history.go
index 055c90a2..3b294790 100644
--- a/internal/tools/sessions_history.go
+++ b/internal/tools/sessions_history.go
@@ -70,12 +70,13 @@ func (t *SessionsHistoryTool) Execute(ctx context.Context, args map[string]any)
 
 	includeTools, _ := args["include_tools"].(bool)
 
-	// Security: validate session belongs to current agent (fail-closed)
-	agentID := resolveAgentIDString(ctx)
-	if agentID == "" {
+	// Security: validate session belongs to current agent (fail-closed).
+	// Session keys use agent_key (e.g. "agent:victoria:..."), not UUID.
+	agentKey := ToolAgentKeyFromCtx(ctx)
+	if agentKey == "" {
 		return ErrorResult("agent context required")
 	}
-	if !strings.HasPrefix(sessionKey, "agent:"+agentID+":") {
+	if !strings.HasPrefix(sessionKey, "agent:"+agentKey+":") {
 		return ErrorResult("access denied: session belongs to a different agent")
 	}
 
diff --git a/internal/tools/sessions_send.go b/internal/tools/sessions_send.go
index abea1edd..51c1c09e 100644
--- a/internal/tools/sessions_send.go
+++ b/internal/tools/sessions_send.go
@@ -69,15 +69,16 @@ func (t *SessionsSendTool) Execute(ctx context.Context, args map[string]any) *Re
 		return ErrorResult("either session_key or label is required")
 	}
 
-	// Security: fail-closed when agentID missing
-	agentID := resolveAgentIDString(ctx)
-	if agentID == "" {
+	// Security: fail-closed when agent key missing.
+	// Session keys use agent_key (e.g. "agent:victoria:..."), not UUID.
+	agentKey := ToolAgentKeyFromCtx(ctx)
+	if agentKey == "" {
 		return ErrorResult("agent context required")
 	}
 
 	// Resolve by label if needed
 	if sessionKey == "" && label != "" {
-		sessions := t.sessions.List(ctx, agentID)
+		sessions := t.sessions.List(ctx, agentKey)
 		for _, s := range sessions {
 			if s.Label == label {
 				sessionKey = s.Key
@@ -90,7 +91,7 @@ func (t *SessionsSendTool) Execute(ctx context.Context, args map[string]any) *Re
 	}
 
 	// Security: validate target session belongs to same agent
-	if !strings.HasPrefix(sessionKey, "agent:"+agentID+":") {
+	if !strings.HasPrefix(sessionKey, "agent:"+agentKey+":") {
 		return ErrorResult("access denied: target session belongs to a different agent")
 	}
 
diff --git a/internal/tools/sessions_test.go b/internal/tools/sessions_test.go
index 3298cd76..537893d6 100644
--- a/internal/tools/sessions_test.go
+++ b/internal/tools/sessions_test.go
@@ -170,7 +170,10 @@ func agentCtx(agentID string) context.Context {
 	if err != nil {
 		return ctx
 	}
-	return store.WithAgentID(ctx, uid)
+	ctx = store.WithAgentID(ctx, uid)
+	// Session keys use agent_key, not UUID.  In production WithToolAgentKey
+	// is set by the agent loop; tests must mirror this.
+	return WithToolAgentKey(ctx, agentID)
 }
 
 func agentCtxWithSandbox(agentID, sandboxKey string) context.Context {