Let any user access anonymous AI messages if they have the ID (#9245)

Author: begelundmuller

runtime/ai/ai.go

@@ -113,10 +113,17 @@ func (r *Runner) Session(ctx context.Context, opts *SessionOptions) (res *Sessio
 		if err != nil {
 			return nil, fmt.Errorf("failed to find session %q: %w", opts.SessionID, err)
 		}
-		retrieveUntilMessageID := session.SharedUntilMessageID
-		if session.OwnerID == opts.Claims.UserID || opts.Claims.SkipChecks {
-			// If the user owns the session or skipCheck enabled, they can see all messages.
-			retrieveUntilMessageID = ""
+
+		// Check access: you can access anonymous sessions, your own sessions, and shared sessions.
+		// For shared sessions, if you are not the owner, you can only see messages up to the SharedUntilMessageID (inclusive).
+		// For sessions without an owner (unauthenticated users using a public project), we don't check access and rely on security by obscurity (generally a decent trade-off, but specifically introduced to get citation links over MCP working for unauthenticated demos).
+		// It's important to respect SkipChecks to ensure access in Rill Developer (where auth is disabled, but SkipChecks is true).
+		var retrieveUntilMessageID string
+		if session.OwnerID != "" && session.OwnerID != opts.Claims.UserID && !opts.Claims.SkipChecks {
+			if session.SharedUntilMessageID == "" {
+				return nil, fmt.Errorf("access denied to session %q", session.ID)
+			}
+			retrieveUntilMessageID = session.SharedUntilMessageID
 		}
 
 		ms, err := catalog.FindAIMessages(ctx, opts.SessionID)
@@ -158,12 +165,6 @@ func (r *Runner) Session(ctx context.Context, opts *SessionOptions) (res *Sessio
 		}
 	}
 
-	// Check access: for now, only allow users to access their own sessions or shared sessions with trimmed messages.
-	// Checking !SkipChecks to ensure access for superusers and for Rill Developer (where auth is disabled and SkipChecks is true).
-	if opts.Claims.UserID != session.OwnerID && !opts.Claims.SkipChecks && session.SharedUntilMessageID == "" {
-		return nil, fmt.Errorf("access denied to session %q", session.ID)
-	}
-
 	// Setup logger
 	logger := r.Runtime.Logger.Named("ai").With(
 		zap.String("instance_id", opts.InstanceID),

runtime/server/chat_test.go

@@ -319,6 +319,74 @@ measures:
 	require.Len(t, get3.Messages, len(res4.Messages))
 }
 
+func TestAnonymousSessionAccess(t *testing.T) {
+	rt, instanceID := testruntime.NewInstance(t)
+	srv, err := server.NewServer(context.Background(), &server.Options{}, rt, zap.NewNop(), ratelimit.NewNoop(), activity.NewNoopClient(), nil)
+	require.NoError(t, err)
+
+	// Claims for an anonymous and authenticated user.
+	anonClaims := &runtime.SecurityClaims{
+		UserID:      "",
+		Permissions: []runtime.Permission{runtime.ReadObjects, runtime.ReadMetrics, runtime.UseAI},
+	}
+	authedClaims := &runtime.SecurityClaims{
+		UserID:      "foo",
+		Permissions: []runtime.Permission{runtime.ReadObjects, runtime.ReadMetrics, runtime.UseAI},
+	}
+
+	// Create an anonymous session with a message directly to avoid LLM calls and keep the test cheap.
+	runner := ai.NewRunner(rt, activity.NewNoopClient())
+	sess, err := runner.Session(t.Context(), &ai.SessionOptions{
+		InstanceID:        instanceID,
+		CreateIfNotExists: true,
+		Claims:            anonClaims,
+		UserAgent:         "test",
+	})
+	require.NoError(t, err)
+	msg, err := sess.CallTool(t.Context(), ai.RoleUser, ai.ListMetricsViewsName, nil, &ai.ListMetricsViewsArgs{})
+	require.NoError(t, err)
+	err = sess.Flush(t.Context())
+	require.NoError(t, err)
+
+	// Anonymous user can access the anonymous session
+	anonCtx := auth.WithClaims(t.Context(), anonClaims)
+	convRes, err := srv.GetConversation(anonCtx, &runtimev1.GetConversationRequest{
+		InstanceId:     instanceID,
+		ConversationId: msg.Call.SessionID,
+	})
+	require.NoError(t, err)
+	require.NotNil(t, convRes.Conversation)
+	require.Equal(t, msg.Call.SessionID, convRes.Conversation.Id)
+
+	// Anonymous user can access the message in the anonymous session
+	msgRes, err := srv.GetAIMessage(anonCtx, &runtimev1.GetAIMessageRequest{
+		InstanceId:     instanceID,
+		ConversationId: msg.Call.SessionID,
+		MessageId:      msg.Call.ID,
+	})
+	require.NoError(t, err)
+	require.Equal(t, msg.Call.ID, msgRes.Message.Id)
+
+	// Authenticated user can also access the anonymous session
+	authedCtx := auth.WithClaims(t.Context(), authedClaims)
+	convRes, err = srv.GetConversation(authedCtx, &runtimev1.GetConversationRequest{
+		InstanceId:     instanceID,
+		ConversationId: msg.Call.SessionID,
+	})
+	require.NoError(t, err)
+	require.NotNil(t, convRes.Conversation)
+	require.Equal(t, msg.Call.SessionID, convRes.Conversation.Id)
+
+	// Authenticated user can access the message in the anonymous session
+	msgRes, err = srv.GetAIMessage(authedCtx, &runtimev1.GetAIMessageRequest{
+		InstanceId:     instanceID,
+		ConversationId: msg.Call.SessionID,
+		MessageId:      msg.Call.ID,
+	})
+	require.NoError(t, err)
+	require.Equal(t, msg.Call.ID, msgRes.Message.Id)
+}
+
 func TestAgentChoiceAndContext(t *testing.T) {
 	// Skip in CI since we make real LLM calls.
 	testmode.Expensive(t)