Repository: google/adk-go
File: artifact/inmemory.go (Go)

Code before:
	"sync"

	"google.golang.org/genai"
	"rsc.io/omap"
	"rsc.io/ordered"
)

// inMemoryService is an in-memory implementation of the Service.
// It is primarily for testing and demonstration purposes.
type inMemoryService struct {
	mu sync.RWMutex
	// ordered(appName, userID, sessionID) -> session
	artifacts omap.Map[string, *genai.Part]
}

func InMemoryService() Service {
	return &inMemoryService{}
}

// fileHasUserNamespace checks if a filename indicates a user scoped artifact.
func fileHasUserNamespace(filename string) bool {
	return strings.HasPrefix(filename, "user:")
}

// userPathPart defines the stirng for the part of the path used by user scope files
const userPathPart = "user"

type artifactKey struct {
	AppName   string
	UserID    string
	SessionID string
	FileName  string
	Version   int64
}

func (ak artifactKey) Encode() string {
	return string(ordered.Encode(ak.AppName, ak.UserID, ak.SessionID, ak.FileName, ordered.Rev(ak.Version)))
}

func (ak *artifactKey) Decode(key string) error {
	var v ordered.Reverse[int64]
	err := ordered.Decode([]byte(key), &ak.AppName, &ak.UserID, &ak.SessionID, &ak.FileName, &v)
	if err != nil {
		return err
	}
	ak.Version = v.Value()
	return nil
}

// scan returns an iterator over all key-value pairs
// in the range begin ≤ key ≤ end.

Review comment (bug):
There's a high chance somebody may create a session with `session_id=user`. Then it will be possible to read artifacts of `user` session from any other by passing `user:key` (although it's still scoped within this one user sessions only so it's not that bad).

Is it viable to have a separate storage for user-scoped artifacts: `inMemoryService.userArtifacts` to avoid any potential bugs.
Or to have some unique key to store user-visible artifacts?


Code after:
	"sync"

	"google.golang.org/genai"
	"rsc.io/omap"
	"rsc.io/ordered"
)

// inMemoryService is an in-memory implementation of the Service.
// It is primarily for testing and demonstration purposes.
type inMemoryService struct {
	mu sync.RWMutex
	// ordered(appName, userID, sessionID) -> session
	artifacts omap.Map[string, *genai.Part]
}

func InMemoryService() Service {
	return &inMemoryService{}
}

// fileHasUserNamespace checks if a filename indicates a user scoped artifact.
func fileHasUserNamespace(filename string) bool {
	return strings.HasPrefix(filename, "user:")
}

// userScopedArtifactKey defines the string for the part of the path used by user scope files
// this key is used to store user artifacts within user namespace (available to all sessions for given app+user_id)
const userScopedArtifactKey = "user"

type artifactKey struct {
	AppName   string
	UserID    string
	SessionID string
	FileName  string
	Version   int64
}

func (ak artifactKey) Encode() string {
	return string(ordered.Encode(ak.AppName, ak.UserID, ak.SessionID, ak.FileName, ordered.Rev(ak.Version)))
}

func (ak *artifactKey) Decode(key string) error {
	var v ordered.Reverse[int64]
	err := ordered.Decode([]byte(key), &ak.AppName, &ak.UserID, &ak.SessionID, &ak.FileName, &v)
	if err != nil {
		return err
	}
	ak.Version = v.Value()
	return nil
}

// scan returns an iterator over all key-value pairs