agents_and_robots/shell/matrix/client.go

// Package matrix wraps mautrix-go for agent use.
package matrix

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"log/slog"
	"os"
	"path/filepath"
	"strings"

	"maunium.net/go/mautrix"
	"maunium.net/go/mautrix/crypto/cryptohelper"
	"maunium.net/go/mautrix/event"
	"maunium.net/go/mautrix/id"

	"github.com/enmanuel/agents/internal/config"
)

// Client wraps a mautrix client with agent-relevant helpers.
type Client struct {
	raw *mautrix.Client
	cfg config.MatrixCfg
}

// New creates and authenticates a Matrix client from config.
// The access token is read from the env var specified in cfg.AccessTokenEnv.
func New(cfg config.MatrixCfg) (*Client, error) {
	token := os.Getenv(cfg.AccessTokenEnv)
	if token == "" {
		return nil, fmt.Errorf("env var %s is not set", cfg.AccessTokenEnv)
	}

	raw, err := mautrix.NewClient(cfg.Homeserver, id.UserID(cfg.UserID), token)
	if err != nil {
		return nil, fmt.Errorf("create matrix client: %w", err)
	}

	if cfg.DeviceID != "" {
		raw.DeviceID = id.DeviceID(cfg.DeviceID)
	}

	return &Client{raw: raw, cfg: cfg}, nil
}

// cryptoIniter abstracts crypto helper creation for testing.
type cryptoIniter interface {
	newHelper(pickleKey []byte, storePath string) (cryptoHelper, error)
}

// cryptoHelper abstracts the mautrix CryptoHelper for testing.
type cryptoHelper interface {
	io.Closer
	Init(ctx context.Context) error
	SetAccountID(id string)
}

// mautrixCryptoIniter is the real implementation using mautrix.
type mautrixCryptoIniter struct {
	raw *mautrix.Client
}

func (m *mautrixCryptoIniter) newHelper(pickleKey []byte, storePath string) (cryptoHelper, error) {
	h, err := cryptohelper.NewCryptoHelper(m.raw, pickleKey, storePath)
	if err != nil {
		return nil, err
	}
	return &mautrixCryptoWrapper{h}, nil
}

type mautrixCryptoWrapper struct {
	*cryptohelper.CryptoHelper
}

func (w *mautrixCryptoWrapper) SetAccountID(id string) { w.DBAccountID = id }

// InitCrypto sets up end-to-end encryption using the mautrix cryptohelper.
// storePath is the SQLite file path for crypto material (e.g. "./agents/<id>/data/crypto/crypto.db").
// pickleKeyHex is a hex-encoded key for encrypting crypto material at rest. If empty,
// falls back to sha256(access_token) for backward compatibility.
// agentID namespaces the crypto state within the database.
// Returns an io.Closer that must be called on agent shutdown to flush the crypto store.
func (c *Client) InitCrypto(ctx context.Context, storePath, pickleKeyHex, agentID string) (io.Closer, error) {
	// Resolve the actual device ID from the server.
	whoami, err := c.raw.Whoami(ctx)
	if err != nil {
		return nil, fmt.Errorf("whoami for crypto init: %w", err)
	}
	c.raw.DeviceID = whoami.DeviceID

	initer := &mautrixCryptoIniter{raw: c.raw}
	closer, helper, err := initCryptoCore(ctx, storePath, pickleKeyHex, c.raw.AccessToken, agentID, initer, slog.Default())
	if err != nil {
		return nil, err
	}

	// Assign the real mautrix crypto helper — this satisfies mautrix.CryptoHelper.
	c.raw.Crypto = helper.(*mautrixCryptoWrapper)
	return closer, nil
}

// initCryptoCore contains the testable logic: pickle key resolution, store
// creation, and auto-recovery on stale crypto.db. Returns (closer, helper, err).
func initCryptoCore(ctx context.Context, storePath, pickleKeyHex, accessToken, agentID string, initer cryptoIniter, logger *slog.Logger) (io.Closer, cryptoHelper, error) {
	pickleKey, err := resolvePickleKey(pickleKeyHex, accessToken)
	if err != nil {
		return nil, nil, err
	}

	if err := os.MkdirAll(filepath.Dir(storePath), 0700); err != nil {
		return nil, nil, fmt.Errorf("create crypto store dir: %w", err)
	}

	helper, err := initHelper(ctx, initer, pickleKey, storePath, agentID)
	if err != nil && strings.Contains(err.Error(), "not marked as shared") {
		logger.Warn("crypto store inconsistent, attempting auto-recovery",
			"store", storePath,
		)
		if removeErr := os.Remove(storePath); removeErr != nil && !os.IsNotExist(removeErr) {
			return nil, nil, fmt.Errorf("auto-recovery: remove stale crypto.db: %w (original: %w)", removeErr, err)
		}
		helper, err = initHelper(ctx, initer, pickleKey, storePath, agentID)
		if err != nil {
			return nil, nil, fmt.Errorf("e2ee init after auto-recovery: %w", err)
		}
		logger.Info("e2ee auto-recovery succeeded")
	} else if err != nil {
		return nil, nil, fmt.Errorf("init e2ee: %w", err)
	}

	return helper, helper, nil
}

func initHelper(ctx context.Context, initer cryptoIniter, pickleKey []byte, storePath, agentID string) (cryptoHelper, error) {
	helper, err := initer.newHelper(pickleKey, storePath)
	if err != nil {
		return nil, fmt.Errorf("create crypto helper: %w", err)
	}
	helper.SetAccountID(agentID)
	if err := helper.Init(ctx); err != nil {
		return nil, err
	}
	return helper, nil
}

// resolvePickleKey decodes a hex key or derives one from the access token.
func resolvePickleKey(pickleKeyHex, accessToken string) ([]byte, error) {
	if pickleKeyHex != "" {
		key, err := hex.DecodeString(pickleKeyHex)
		if err != nil {
			return nil, fmt.Errorf("decode pickle_key_env: %w", err)
		}
		return key, nil
	}
	sum := sha256.Sum256([]byte(accessToken))
	return sum[:], nil
}

// SendText sends a plain-text message to a room.
// If the room has E2EE enabled and crypto is initialized, the message is encrypted automatically.
func (c *Client) SendText(ctx context.Context, roomID, text string) error {
	_, err := c.raw.SendText(ctx, id.RoomID(roomID), text)
	return err
}

// SendMarkdown sends a formatted (Markdown) message to a room.
func (c *Client) SendMarkdown(ctx context.Context, roomID, markdown string) error {
	content := event.MessageEventContent{
		MsgType:       event.MsgText,
		Body:          markdown,
		Format:        event.FormatHTML,
		FormattedBody: markdown, // mautrix can render markdown -> HTML if needed
	}
	_, err := c.raw.SendMessageEvent(ctx, id.RoomID(roomID), event.EventMessage, content)
	return err
}

// SendReaction sends a reaction to an event.
func (c *Client) SendReaction(ctx context.Context, roomID, eventID, reaction string) error {
	_, err := c.raw.SendReaction(ctx, id.RoomID(roomID), id.EventID(eventID), reaction)
	return err
}

// SendTyping sets the typing indicator in a room.
func (c *Client) SendTyping(ctx context.Context, roomID string, typing bool) error {
	_, err := c.raw.UserTyping(ctx, id.RoomID(roomID), typing, 5000)
	return err
}

// Raw returns the underlying mautrix.Client for advanced use.
func (c *Client) Raw() *mautrix.Client {
	return c.raw
}