Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
109 changes: 109 additions & 0 deletions admin/server/deployment_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,8 @@ func TestDeploymentJWTs(t *testing.T) {
require.Equal(t, u1.ID, c.Subject)
require.Equal(t, u1.Email, c.Attrs["email"])
require.Equal(t, true, c.Attrs["admin"])
require.NotContains(t, c.Attrs, "service")
require.NotContains(t, c.Attrs, "public_url")

perms := c.Instances[depl.RuntimeInstanceID]
requireHasPerms(t, perms, runtime.ReadObjects, runtime.ReadMetrics, runtime.ReadAPI, runtime.UseAI)
Expand Down Expand Up @@ -120,6 +122,7 @@ func TestDeploymentJWTs(t *testing.T) {
require.NotEqual(t, u1.ID, c.Subject)
require.Equal(t, "embed@example.com", c.Attrs["email"])
require.Equal(t, "foo", c.Attrs["tenant_id"])
require.Equal(t, true, c.Attrs["external_user"])
require.NotContains(t, c.Attrs, "admin")

// Subject hash must be deterministic across calls with the same external id and project.
Expand All @@ -131,6 +134,112 @@ func TestDeploymentJWTs(t *testing.T) {
require.NoError(t, err)
require.Equal(t, c.Subject, decodeJWT(t, resp2.AccessToken).Subject)
})

t.Run("Service token", func(t *testing.T) {
svc, err := u1Client.CreateService(t.Context(), &adminv1.CreateServiceRequest{
Org: "jwt-test",
Name: "embed-service",
Project: "jwt-project",
ProjectRoleName: "admin",
})
require.NoError(t, err)
tok, err := u1Client.IssueServiceAuthToken(t.Context(), &adminv1.IssueServiceAuthTokenRequest{
Org: "jwt-test",
ServiceName: "embed-service",
})
require.NoError(t, err)

svcClient := adm.NewClient(t, tok.Token)

// The service's own token has the service ID as subject, shared by every caller, and is marked as a service.
proj, err := svcClient.GetProject(t.Context(), &adminv1.GetProjectRequest{
Org: "jwt-test",
Project: "jwt-project",
})
require.NoError(t, err)
c := decodeJWT(t, proj.Jwt)
require.Equal(t, svc.Service.Id, c.Subject)
require.Equal(t, true, c.Attrs["service"])
require.NotContains(t, c.Attrs, "embed")
require.NotContains(t, c.Attrs, "external_user")
requireHasPerms(t, c.Instances[depl.RuntimeInstanceID], runtime.UseAI)

// An embed without an external user ID keeps the service ID as subject for every viewer, and nothing marks the token as identifying one viewer.
resp, err := svcClient.GetIFrame(t.Context(), &adminv1.GetIFrameRequest{
Org: "jwt-test",
Project: "jwt-project",
Resource: "foo",
})
require.NoError(t, err)
c = decodeJWT(t, resp.AccessToken)
require.Equal(t, svc.Service.Id, c.Subject)
require.Equal(t, true, c.Attrs["embed"])
require.Equal(t, true, c.Attrs["service"])
require.NotContains(t, c.Attrs, "external_user")
require.NotContains(t, c.Attrs, "public_url")
requireHasPerms(t, c.Instances[depl.RuntimeInstanceID], runtime.UseAI)

// With an external user ID, each viewer gets their own subject and the token says so. The service itself is no longer the principal.
resp, err = svcClient.GetIFrame(t.Context(), &adminv1.GetIFrameRequest{
Org: "jwt-test",
Project: "jwt-project",
Resource: "foo",
ExternalUserId: "external-customer-42",
})
require.NoError(t, err)
c = decodeJWT(t, resp.AccessToken)
require.True(t, strings.HasPrefix(c.Subject, "ext_"))
require.Equal(t, true, c.Attrs["embed"])
require.Equal(t, true, c.Attrs["external_user"])
require.NotContains(t, c.Attrs, "service")

// The markers are derived by the server: custom attributes cannot claim a per-viewer subject or hide the shared one.
spoofed, err := structpb.NewStruct(map[string]any{
"external_user": true,
"service": false,
"public_url": false,
"tenant_id": "foo",
})
require.NoError(t, err)
resp, err = svcClient.GetIFrame(t.Context(), &adminv1.GetIFrameRequest{
Org: "jwt-test",
Project: "jwt-project",
Resource: "foo",
For: &adminv1.GetIFrameRequest_Attributes{Attributes: spoofed},
})
require.NoError(t, err)
c = decodeJWT(t, resp.AccessToken)
require.Empty(t, c.Subject)
require.Equal(t, "foo", c.Attrs["tenant_id"])
require.Equal(t, true, c.Attrs["embed"])
require.NotContains(t, c.Attrs, "external_user")
require.NotContains(t, c.Attrs, "service")
require.NotContains(t, c.Attrs, "public_url")
})

t.Run("GetProject for public URL", func(t *testing.T) {
tkn, err := u1Client.IssueMagicAuthToken(t.Context(), &adminv1.IssueMagicAuthTokenRequest{
Org: "jwt-test",
Project: "jwt-project",
Resources: []*adminv1.ResourceName{{Type: runtime.ResourceKindExplore, Name: "foo"}},
})
require.NoError(t, err)

resp, err := adm.NewClient(t, tkn.Token).GetProject(t.Context(), &adminv1.GetProjectRequest{
Org: "jwt-test",
Project: "jwt-project",
})
require.NoError(t, err)

// Everyone with the link shares the magic token ID as subject and the creator's attributes, so the token is marked as a public URL.
c := decodeJWT(t, resp.Jwt)
require.NotEmpty(t, c.Subject)
require.NotEqual(t, u1.ID, c.Subject)
require.Equal(t, u1.Email, c.Attrs["email"])
require.Equal(t, true, c.Attrs["public_url"])
require.NotContains(t, c.Attrs, "embed")
requireHasPerms(t, c.Instances[depl.RuntimeInstanceID], runtime.UseAI)
})
}

// jwtPayload mirrors the subset of the runtime JWT payload that the assertions touch.
Expand Down
27 changes: 27 additions & 0 deletions admin/server/runtime_jwt.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package server

import (
"context"
"maps"
"time"

"github.com/rilldata/rill/admin/database"
Expand Down Expand Up @@ -85,6 +86,7 @@ func (s *Server) issueRuntimeToken(ctx context.Context, opts *issueRuntimeTokenO
var subject string
var attr map[string]any
var rules []*runtimev1.SecurityRule
var isService, isPublicURL bool
switch {
case opts.forOwner:
if opts.externalUserID != "" {
Expand All @@ -105,6 +107,7 @@ func (s *Server) issueRuntimeToken(ctx context.Context, opts *issueRuntimeTokenO
rules = append(rules, securityRulesFromResources(restrict, resources)...)
case auth.OwnerTypeService:
subject = claims.OwnerID()
isService = true
a, err := s.jwtAttributesForService(ctx, claims.OwnerID(), opts.projectPermissions)
if err != nil {
return "", err
Expand All @@ -116,6 +119,7 @@ func (s *Server) issueRuntimeToken(ctx context.Context, opts *issueRuntimeTokenO
if !ok {
return "", status.Errorf(codes.Internal, "unexpected type %T for magic auth token model", claims.AuthTokenModel())
}
isPublicURL = true
attr = mdl.Attributes
magicRules, err := securityRulesFromMagicAuthToken(mdl)
if err != nil {
Expand Down Expand Up @@ -156,6 +160,29 @@ func (s *Server) issueRuntimeToken(ctx context.Context, opts *issueRuntimeTokenO
subject = subjectForExternalUser(opts.externalUserID, opts.project.ID)
}

// Mark tokens whose subject does not identify one viewer, so the runtime can keep per-user features such as AI memory off for them (see runtime.ResolveFeatureFlags):
// a service is not a person, a public URL puts the magic token ID on everyone with the link, and an embed (marked by extraAttributes) serves many viewers.
// An external_user_id gives each viewer their own subject, which external_user records.
// The markers are derived here only: a value in stored or caller-supplied attributes is dropped, so attributes such as external_user: true cannot turn memory on for a shared subject.
attr = maps.Clone(attr)
delete(attr, "service")
delete(attr, "public_url")
delete(attr, "external_user")
if isService || isPublicURL || opts.externalUserID != "" {
if attr == nil {
attr = make(map[string]any, 1)
}
if isService {
attr["service"] = true
}
if isPublicURL {
attr["public_url"] = true
}
if opts.externalUserID != "" {
attr["external_user"] = true
}
}

// overrideResources replaces any principal-derived resource rules with a fresh allow-list.
if len(opts.overrideResources) > 0 {
rules = securityRulesFromResources(true, opts.overrideResources)
Expand Down
2 changes: 1 addition & 1 deletion docs/docs/developers/build/ai-configuration.md
Original file line number Diff line number Diff line change
Expand Up @@ -194,7 +194,7 @@ In addition to the project-wide `ai_instructions` you write, [AI Chat](/guide/ai

If several users keep teaching the AI the same thing, add it to `ai_instructions` so everyone benefits. Memory never stores data values, secrets, or anything `ai_instructions` already say.

Memory is on by default for logged-in users. To turn it off for a project, disable the `chat_memory` feature flag in `rill.yaml`:
Memory is on by default for logged-in users and for embed or service token viewers identified with `external_user_id`. It is never on for public URL viewers, or for service tokens and embeds without `external_user_id`, whatever the flag says, because everyone using the token would share one memory. To turn it off for a project, disable the `chat_memory` feature flag in `rill.yaml`:

```yaml
features:
Expand Down
2 changes: 2 additions & 0 deletions docs/docs/developers/embed/iframe.md
Original file line number Diff line number Diff line change
Expand Up @@ -328,6 +328,8 @@ Only one of `user_email` or `attributes` can be provided for a given iframe. The
- `attributes`: Passes the provided attributes through directly. Make sure to include all attributes referenced in your security policies (e.g. `email`, `domain`, `admin`, or custom attributes like `tenant_id`). Does not enable per-user state on its own; combine with `external_user_id` to enable per-user state.
- `external_user_id`: Any stable identifier for the end user. This is usually the user's ID in your own database. Setting it enables per-user state such as AI chat history.

[AI memory](/guide/ai/ai-chat#memory) is per-user state too: it is only available in an embed with `external_user_id`. Without it, every viewer of the embed would share one memory, so memory is off.

The same attributes decide what an embedded canvas shows. Canvas rows, tabs and components can have [`if` conditions](/developers/build/dashboards/canvas#show-content-to-some-viewers-only) on them, such as `'{{ eq .user.plan "premium" }}'` to show premium content only to tenants on that plan. Every embed also carries the attribute `embed: true`, so a condition like `'{{ not .user.embed }}'` keeps content out of embedded dashboards.

## Embedding the project vs embedding an individual dashboard
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -203,3 +203,5 @@ Only one of `user_email` or `attributes` can be provided for a given token. The
- `user_email`: Looks up the user in Rill Cloud by email and populates their standard attributes. If no matching user is found, it generates limited attributes with only the fields `email`, `domain` and `admin` (set to `false`). Does not enable per-user state on its own; combine with `external_user_id` to enable per-user state.
- `attributes`: Passes the provided attributes through directly. Make sure to include all attributes referenced in your security policies (e.g. `email`, `domain`, `admin`, or custom attributes like `tenant_id`). Does not enable per-user state on its own; combine with `external_user_id` to enable per-user state.
- `external_user_id`: Any stable identifier for an external end user. This is usually the user's ID in your own database. Setting it enables per-user state such as AI chat history. Can be combined with `user_email` or `attributes`.

[AI memory](/guide/ai/ai-chat#memory) is per-user state too: it is only available with `external_user_id`. Without it, everyone using the token would share one memory, so memory is off.
2 changes: 1 addition & 1 deletion docs/docs/guide/ai/ai-chat.md
Original file line number Diff line number Diff line change
Expand Up @@ -172,7 +172,7 @@ Open **Manage memory** from the chat sidebar (or **Settings > AI memory** in Ril

Turn **Use memory** off to keep your memories but stop the AI from using or updating them until you turn it back on.

Memory is available in Rill Cloud for logged-in users and in Rill Developer, where it is stored locally and never committed to git. It is not used by AI reports, shared conversations viewed by other users, or external AI clients connected through MCP. Project administrators can turn memory off for a project with the `chat_memory` feature flag; see [AI Configuration](/developers/build/ai-configuration#user-memory).
Memory is per user. It is available in Rill Cloud for logged-in users and in Rill Developer, where it is stored locally and never committed to git. In [embedded dashboards](/developers/embed/iframe) and with [service tokens](/guide/administration/access-tokens/service-tokens) it is only available when the viewer is identified with `external_user_id`; otherwise everyone using the token would share one memory, so it is off. It is never available through [public URLs](/guide/dashboards/public-urls), and it is not used by AI reports, shared conversations viewed by other users, or external AI clients connected through MCP. Project administrators can turn memory off for a project with the `chat_memory` feature flag; see [AI Configuration](/developers/build/ai-configuration#user-memory).

## Use in Your Favorite AI Client

Expand Down
7 changes: 5 additions & 2 deletions runtime/ai/memory.go
Original file line number Diff line number Diff line change
Expand Up @@ -58,12 +58,15 @@ var (
)

// MemoryAvailable reports whether memory is available to the given claims in an instance.
// It is false for users without the AI permission, for anonymous users, and when the chat_memory feature flag is off.
// Memory is keyed on the user ID, so it is only available when that ID identifies one viewer:
// it is false for users without the AI permission, for anonymous users, and when the chat_memory feature flag is off.
// The flag is also resolved to false for tokens whose subject does not identify one viewer (see runtime.ResolveFeatureFlags):
// public URLs, and service tokens or embeds without an external_user_id.
func MemoryAvailable(ctx context.Context, rt *runtime.Runtime, instanceID string, claims *runtime.SecurityClaims) (bool, error) {
if claims == nil || !claims.Can(runtime.UseAI) {
return false, nil
}
// Anonymous users on Rill Cloud (public projects, magic tokens) have no stable identity to key memory on.
// Anonymous users on Rill Cloud (public projects) have no identity to key memory on.
// In Rill Developer, auth is disabled so SkipChecks is true and the empty user ID acts as the single local user.
if claims.UserID == "" && !claims.SkipChecks {
return false, nil
Expand Down
6 changes: 6 additions & 0 deletions runtime/ai/memory_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -213,6 +213,12 @@ func TestMemoryGates(t *testing.T) {
{"no ai permission", "rill/0.1", &runtime.SecurityClaims{UserID: uuid.NewString()}, false},
{"report", "rill/report", &runtime.SecurityClaims{UserID: uuid.NewString(), SkipChecks: true}, false},
{"mcp client", "claude-desktop/1.0", &runtime.SecurityClaims{UserID: uuid.NewString(), SkipChecks: true}, false},
// A service token puts the service ID on every caller, and an embed without an external_user_id puts the issuer's identity on every viewer; with one, each viewer has their own subject.
{"service token", "rill/0.1", &runtime.SecurityClaims{UserID: uuid.NewString(), UserAttributes: map[string]any{"name": "my-service", "admin": true, "service": true}, Permissions: []runtime.Permission{runtime.UseAI}}, false},
{"service token embed", "rill/0.1", &runtime.SecurityClaims{UserID: uuid.NewString(), UserAttributes: map[string]any{"name": "my-service", "admin": true, "service": true, "embed": true}, Permissions: []runtime.Permission{runtime.UseAI}}, false},
{"external user embed", "rill/0.1", &runtime.SecurityClaims{UserID: "ext_" + uuid.NewString(), UserAttributes: map[string]any{"embed": true, "external_user": true}, Permissions: []runtime.Permission{runtime.UseAI}}, true},
// A public URL token carries the magic auth token ID as subject and the creator's attributes, shared by everyone with the link.
{"public URL", "rill/0.1", &runtime.SecurityClaims{UserID: uuid.NewString(), UserAttributes: map[string]any{"email": "creator@example.com", "domain": "example.com", "admin": true, "public_url": true}, Permissions: []runtime.Permission{runtime.UseAI}}, false},
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
Expand Down
17 changes: 16 additions & 1 deletion runtime/feature_flags.go
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,8 @@ var defaultFeatureFlags = map[string]string{
"developer_chat": "true",
// Controls whether charts are rendered in AI chats
"chat_charts": "true",
// Controls whether the AI chat learns and uses per-user memories (preferences and clarifications remembered across conversations)
// Controls whether the AI chat learns and uses per-user memories (preferences and clarifications remembered across conversations).
// Always resolved to false when the token's subject does not identify one viewer: public URLs, and service tokens or embeds without an external_user_id.
"chat_memory": "true",
// Controls whether to show/hide deploy related actions.
"deploy": "true",
Expand Down Expand Up @@ -150,5 +151,19 @@ func ResolveFeatureFlags(inst *drivers.Instance, userAttributes map[string]any,
featureFlags[chatMemoryKey] = false
}

// Memory is keyed on the token's subject, so it is off whenever the subject does not identify one viewer, whatever the project configures:
// a public URL puts the magic token ID on everyone with the link, a service token puts the service ID on every caller,
// and an embed without an external_user_id puts the issuer's identity on every viewer's token.
// An external_user_id gives each viewer their own subject, so a service token or embed with one keeps memory.
// The admin server sets these attributes when it issues the token.
// The runtime enforces this flag (see ai.MemoryAvailable), and the UI hides the memory controls when it is off.
embed, _ := userAttributes["embed"].(bool)
service, _ := userAttributes["service"].(bool)
externalUser, _ := userAttributes["external_user"].(bool)
publicURL, _ := userAttributes["public_url"].(bool)
if publicURL || ((embed || service) && !externalUser) {
featureFlags[chatMemoryKey] = false
}

return featureFlags, nil
}
37 changes: 37 additions & 0 deletions runtime/feature_flags_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -194,3 +194,40 @@ func Test_ResolveFeatureFlags_CloudEditingDisabledAnnotation(t *testing.T) {
})
}
}

// Test_ResolveFeatureFlags_ChatMemoryPrincipals checks that chat_memory is off for tokens whose subject does not identify one viewer, whatever the project configures:
// public URLs, and service tokens or embeds without an external_user_id.
func Test_ResolveFeatureFlags_ChatMemoryPrincipals(t *testing.T) {
inst := &drivers.Instance{
FeatureFlags: map[string]string{
"chat": "true",
"dashboard_chat": "true",
"chat_memory": "true",
},
}

tests := []struct {
name string
userAttrs map[string]any
chatMemory bool
}{
{"rill user", map[string]any{"email": "user@example.com"}, true},
{"anonymous", nil, true},
{"embed without external user", map[string]any{"embed": true, "tenant_id": "foo"}, false},
{"embed with external user", map[string]any{"embed": true, "external_user": true, "tenant_id": "foo"}, true},
{"ephemeral token with external user", map[string]any{"external_user": true}, true},
{"service token", map[string]any{"name": "my-service", "admin": true, "service": true}, false},
{"service token embed with external user", map[string]any{"name": "my-service", "admin": true, "service": true, "embed": true, "external_user": true}, true},
{"public URL", map[string]any{"email": "creator@example.com", "public_url": true}, false},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
featureFlags, err := ResolveFeatureFlags(inst, test.userAttrs, true)
require.NoError(t, err)
require.Equal(t, test.chatMemory, featureFlags["chatMemory"])
// Only memory is affected: chat itself stays as configured.
require.True(t, featureFlags["chat"])
require.True(t, featureFlags["dashboardChat"])
})
}
}
2 changes: 1 addition & 1 deletion runtime/server/ai_memory.go
Original file line number Diff line number Diff line change
Expand Up @@ -188,7 +188,7 @@ func (s *Server) UpdateAIMemorySettings(ctx context.Context, req *runtimev1.Upda
}

// checkAIMemoryAccess checks that the caller may use AI and that memory is available to them.
// It returns a FailedPrecondition error when memory is not available: for anonymous users on Rill Cloud and when the chat_memory feature flag is off.
// It returns a FailedPrecondition error when memory is not available: for anonymous users on Rill Cloud, for public URL viewers and for service tokens or embeds without an external_user_id, and when the chat_memory feature flag is off.
func (s *Server) checkAIMemoryAccess(ctx context.Context, instanceID string) (*runtime.SecurityClaims, error) {
claims := auth.GetClaims(ctx, instanceID)
if !claims.Can(runtime.UseAI) {
Expand Down
Loading
Loading