tatara-cli v0.1.0 + tatara-memory v0.1.3 Implementation Plan¶

r.Route("/v1", func(r chi.Router) { ... })

// Before:
r.Route("/v1", func(v1 chi.Router) {
    v1.Post("/memories", ...)
    v1.Get("/memories/{id}", ...)
    ...
})

// After:
r.Post("/memories", ...)
r.Get("/memories/{id}", ...)
...

// Package version holds build-time version info populated via ldflags.
package version

var (
    Version = "dev"
    Commit  = "unknown"
    Date    = "unknown"
)

package main

import (
    "fmt"
    "os"

    "github.com/szymonrychu/tatara-cli/internal/version"
)

func main() {
    fmt.Printf("tatara %s (%s) built %s\n", version.Version, version.Commit, version.Date)
    os.Exit(0)
}

package auth

import (
    "encoding/json"
    "errors"
    "fmt"
    "os"
    "path/filepath"
    "time"
)

type Token struct {
    AccessToken  string    `json:"access_token"`
    RefreshToken string    `json:"refresh_token"`
    IDToken      string    `json:"id_token,omitempty"`
    ExpiresAt    time.Time `json:"expires_at"`
    TokenType    string    `json:"token_type"`
}

var ErrNoToken = errors.New("auth: no token (run `tatara login`)")

func DefaultTokenPath() (string, error) {
    dir := os.Getenv("XDG_CONFIG_HOME")
    if dir == "" {
        h, err := os.UserHomeDir()
        if err != nil {
            return "", fmt.Errorf("auth: resolve home: %w", err)
        }
        dir = filepath.Join(h, ".config")
    }
    return filepath.Join(dir, "tatara", "token.json"), nil
}

func LoadToken(path string) (*Token, error) {
    b, err := os.ReadFile(path)
    if errors.Is(err, os.ErrNotExist) {
        return nil, ErrNoToken
    }
    if err != nil {
        return nil, fmt.Errorf("auth: read token: %w", err)
    }
    var t Token
    if err := json.Unmarshal(b, &t); err != nil {
        return nil, fmt.Errorf("auth: parse token: %w", err)
    }
    return &t, nil
}

func SaveToken(path string, t *Token) error {
    if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
        return fmt.Errorf("auth: mkdir: %w", err)
    }
    b, err := json.MarshalIndent(t, "", "  ")
    if err != nil {
        return fmt.Errorf("auth: marshal token: %w", err)
    }
    f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o600)
    if err != nil {
        return fmt.Errorf("auth: open token file: %w", err)
    }
    defer f.Close()
    if _, err := f.Write(b); err != nil {
        return fmt.Errorf("auth: write token: %w", err)
    }
    return nil
}

func DeleteToken(path string) error {
    err := os.Remove(path)
    if errors.Is(err, os.ErrNotExist) {
        return nil
    }
    return err
}

package auth_test

import (
    "os"
    "path/filepath"
    "testing"
    "time"

    "github.com/stretchr/testify/require"

    "github.com/szymonrychu/tatara-cli/internal/auth"
)

func TestTokenRoundTrip(t *testing.T) {
    dir := t.TempDir()
    path := filepath.Join(dir, "token.json")

    in := &auth.Token{
        AccessToken:  "at",
        RefreshToken: "rt",
        ExpiresAt:    time.Now().Add(time.Hour).UTC().Truncate(time.Second),
        TokenType:    "Bearer",
    }
    require.NoError(t, auth.SaveToken(path, in))

    info, err := os.Stat(path)
    require.NoError(t, err)
    require.Equal(t, os.FileMode(0o600), info.Mode().Perm())

    got, err := auth.LoadToken(path)
    require.NoError(t, err)
    require.Equal(t, in.AccessToken, got.AccessToken)
    require.True(t, in.ExpiresAt.Equal(got.ExpiresAt))
}

func TestLoadTokenMissing(t *testing.T) {
    _, err := auth.LoadToken(filepath.Join(t.TempDir(), "absent.json"))
    require.ErrorIs(t, err, auth.ErrNoToken)
}

func TestDefaultTokenPathHonorsXDG(t *testing.T) {
    t.Setenv("XDG_CONFIG_HOME", "/tmp/xdg-test")
    p, err := auth.DefaultTokenPath()
    require.NoError(t, err)
    require.Equal(t, "/tmp/xdg-test/tatara/token.json", p)
}

package auth

import (
    "fmt"
    "os"
    "syscall"
)

type FileLock struct{ f *os.File }

func AcquireLock(path string) (*FileLock, error) {
    f, err := os.OpenFile(path+".lock", os.O_CREATE|os.O_RDWR, 0o600)
    if err != nil {
        return nil, fmt.Errorf("auth: open lock: %w", err)
    }
    if err := syscall.Flock(int(f.Fd()), syscall.LOCK_EX); err != nil {
        _ = f.Close()
        return nil, fmt.Errorf("auth: flock: %w", err)
    }
    return &FileLock{f: f}, nil
}

func (l *FileLock) Release() error {
    if err := syscall.Flock(int(l.f.Fd()), syscall.LOCK_UN); err != nil {
        return err
    }
    return l.f.Close()
}

package auth

import (
    "context"
    "encoding/json"
    "errors"
    "fmt"
    "io"
    "net/http"
    "net/url"
    "strings"
    "time"
)

type DeviceCode struct {
    DeviceCode              string `json:"device_code"`
    UserCode                string `json:"user_code"`
    VerificationURI         string `json:"verification_uri"`
    VerificationURIComplete string `json:"verification_uri_complete"`
    ExpiresIn               int    `json:"expires_in"`
    Interval                int    `json:"interval"`
}

type DeviceFlow struct {
    HTTP          *http.Client
    Issuer        string
    ClientID      string
    Scope         string
}

func (d *DeviceFlow) Start(ctx context.Context) (*DeviceCode, error) {
    body := url.Values{}
    body.Set("client_id", d.ClientID)
    body.Set("scope", d.Scope)
    req, err := http.NewRequestWithContext(ctx, http.MethodPost,
        d.Issuer+"/protocol/openid-connect/auth/device", strings.NewReader(body.Encode()))
    if err != nil {
        return nil, err
    }
    req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
    req.Header.Set("Accept", "application/json")
    resp, err := d.client().Do(req)
    if err != nil {
        return nil, fmt.Errorf("auth: device start: %w", err)
    }
    defer resp.Body.Close()
    if resp.StatusCode != http.StatusOK {
        b, _ := io.ReadAll(resp.Body)
        return nil, fmt.Errorf("auth: device start %d: %s", resp.StatusCode, string(b))
    }
    var dc DeviceCode
    if err := json.NewDecoder(resp.Body).Decode(&dc); err != nil {
        return nil, fmt.Errorf("auth: parse device code: %w", err)
    }
    return &dc, nil
}

func (d *DeviceFlow) Poll(ctx context.Context, code *DeviceCode) (*Token, error) {
    interval := time.Duration(code.Interval) * time.Second
    if interval == 0 {
        interval = 5 * time.Second
    }
    deadline := time.Now().Add(time.Duration(code.ExpiresIn) * time.Second)
    for {
        select {
        case <-ctx.Done():
            return nil, ctx.Err()
        default:
        }
        tok, err := d.exchange(ctx, code.DeviceCode)
        if err == nil {
            return tok, nil
        }
        if errors.Is(err, ErrAuthPending) {
            // wait and retry
        } else if errors.Is(err, ErrSlowDown) {
            interval += 5 * time.Second
        } else {
            return nil, err
        }
        if time.Now().After(deadline) {
            return nil, fmt.Errorf("auth: device code expired")
        }
        select {
        case <-ctx.Done():
            return nil, ctx.Err()
        case <-time.After(interval):
        }
    }
}

var (
    ErrAuthPending = errors.New("authorization_pending")
    ErrSlowDown    = errors.New("slow_down")
    ErrAccessDenied = errors.New("access_denied")
)

func (d *DeviceFlow) exchange(ctx context.Context, deviceCode string) (*Token, error) {
    body := url.Values{}
    body.Set("grant_type", "urn:ietf:params:oauth:grant-type:device_code")
    body.Set("device_code", deviceCode)
    body.Set("client_id", d.ClientID)
    req, _ := http.NewRequestWithContext(ctx, http.MethodPost,
        d.Issuer+"/protocol/openid-connect/token", strings.NewReader(body.Encode()))
    req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
    resp, err := d.client().Do(req)
    if err != nil {
        return nil, fmt.Errorf("auth: token exchange: %w", err)
    }
    defer resp.Body.Close()
    raw, _ := io.ReadAll(resp.Body)
    if resp.StatusCode == http.StatusOK {
        var tr struct {
            AccessToken  string `json:"access_token"`
            RefreshToken string `json:"refresh_token"`
            IDToken      string `json:"id_token"`
            ExpiresIn    int    `json:"expires_in"`
            TokenType    string `json:"token_type"`
        }
        if err := json.Unmarshal(raw, &tr); err != nil {
            return nil, fmt.Errorf("auth: parse token: %w", err)
        }
        return &Token{
            AccessToken:  tr.AccessToken,
            RefreshToken: tr.RefreshToken,
            IDToken:      tr.IDToken,
            ExpiresAt:    time.Now().Add(time.Duration(tr.ExpiresIn) * time.Second),
            TokenType:    tr.TokenType,
        }, nil
    }
    var e struct {
        Error string `json:"error"`
    }
    _ = json.Unmarshal(raw, &e)
    switch e.Error {
    case "authorization_pending":
        return nil, ErrAuthPending
    case "slow_down":
        return nil, ErrSlowDown
    case "access_denied":
        return nil, ErrAccessDenied
    default:
        return nil, fmt.Errorf("auth: %s: %s", e.Error, string(raw))
    }
}

func (d *DeviceFlow) client() *http.Client {
    if d.HTTP != nil {
        return d.HTTP
    }
    return &http.Client{Timeout: 30 * time.Second}
}

func RefreshToken(ctx context.Context, issuer, clientID string, t *Token) (*Token, error)

package client

import (
    "fmt"
    "os"
    "path/filepath"

    "gopkg.in/yaml.v3"
)

const DefaultBaseURL = "https://tatara.szymonrichert.pl/api/v1/memory"

type FileConfig struct {
    BaseURL string `yaml:"baseUrl"`
    Issuer  string `yaml:"issuer"`
}

// ResolveBaseURL returns the first non-empty value of: flag, env, file, default.
func ResolveBaseURL(flag, env string, file *FileConfig) string {
    if flag != "" {
        return flag
    }
    if env != "" {
        return env
    }
    if file != nil && file.BaseURL != "" {
        return file.BaseURL
    }
    return DefaultBaseURL
}

func DefaultConfigPath() (string, error) {
    dir := os.Getenv("XDG_CONFIG_HOME")
    if dir == "" {
        h, err := os.UserHomeDir()
        if err != nil {
            return "", fmt.Errorf("client: home: %w", err)
        }
        dir = filepath.Join(h, ".config")
    }
    return filepath.Join(dir, "tatara", "config.yaml"), nil
}

func LoadConfig(path string) (*FileConfig, error) {
    b, err := os.ReadFile(path)
    if os.IsNotExist(err) {
        return &FileConfig{}, nil
    }
    if err != nil {
        return nil, err
    }
    var c FileConfig
    if err := yaml.Unmarshal(b, &c); err != nil {
        return nil, err
    }
    return &c, nil
}

package client

import (
    "bytes"
    "context"
    "encoding/json"
    "errors"
    "fmt"
    "io"
    "net/http"
    "time"

    "github.com/szymonrychu/tatara-cli/internal/auth"
)

type RefreshFunc func(ctx context.Context, t *auth.Token) (*auth.Token, error)

type Client struct {
    base    string
    http    *http.Client
    token   *auth.Token
    refresh RefreshFunc
    save    func(*auth.Token) error
}

type Config struct {
    BaseURL string
    Token   *auth.Token
    Refresh RefreshFunc
    Save    func(*auth.Token) error
    HTTP    *http.Client
}

func New(cfg Config) (*Client, error) {
    if cfg.BaseURL == "" {
        return nil, errors.New("client: BaseURL required")
    }
    h := cfg.HTTP
    if h == nil {
        h = &http.Client{Timeout: 60 * time.Second}
    }
    return &Client{
        base:    cfg.BaseURL,
        http:    h,
        token:   cfg.Token,
        refresh: cfg.Refresh,
        save:    cfg.Save,
    }, nil
}

// Do issues an authenticated request. If body is non-nil, it must be JSON-encodable.
func (c *Client) Do(ctx context.Context, method, path string, body any) (*http.Response, error) {
    if err := c.ensureFresh(ctx); err != nil {
        return nil, err
    }
    var bodyReader io.Reader
    if body != nil {
        switch b := body.(type) {
        case []byte:
            bodyReader = bytes.NewReader(b)
        case io.Reader:
            bodyReader = b
        default:
            buf := &bytes.Buffer{}
            if err := json.NewEncoder(buf).Encode(body); err != nil {
                return nil, fmt.Errorf("client: encode body: %w", err)
            }
            bodyReader = buf
        }
    }
    req, err := http.NewRequestWithContext(ctx, method, c.base+path, bodyReader)
    if err != nil {
        return nil, fmt.Errorf("client: build request: %w", err)
    }
    req.Header.Set("Accept", "application/json")
    if bodyReader != nil {
        req.Header.Set("Content-Type", "application/json")
    }
    if c.token != nil {
        req.Header.Set("Authorization", "Bearer "+c.token.AccessToken)
    }
    return c.http.Do(req)
}

func (c *Client) ensureFresh(ctx context.Context) error {
    if c.token == nil {
        return auth.ErrNoToken
    }
    if time.Until(c.token.ExpiresAt) > 30*time.Second {
        return nil
    }
    if c.refresh == nil {
        return nil
    }
    nt, err := c.refresh(ctx, c.token)
    if err != nil {
        return fmt.Errorf("client: refresh: %w", err)
    }
    c.token = nt
    if c.save != nil {
        return c.save(nt)
    }
    return nil
}

package mcp

import (
    "context"
    "encoding/json"
    "fmt"
    "io"
    "net/http"
    "strings"

    "github.com/szymonrychu/tatara-cli/internal/client"
)

type Tool struct {
    Name        string
    Description string
    Schema      json.RawMessage // JSON schema for arguments
    Build       func(args map[string]any) (method, path string, body any, err error)
}

func AllTools() []Tool {
    return []Tool{
        {
            Name:        "create_memory",
            Description: "Insert a new text memory. Returns the track_id.",
            Schema:      json.RawMessage(`{"type":"object","properties":{"text":{"type":"string"},"metadata":{"type":"object","additionalProperties":{"type":"string"}}},"required":["text"]}`),
            Build: func(a map[string]any) (string, string, any, error) {
                return http.MethodPost, "/memories", a, nil
            },
        },
        {
            Name:        "get_memory",
            Description: "Retrieve a memory by track_id.",
            Schema:      json.RawMessage(`{"type":"object","properties":{"id":{"type":"string"}},"required":["id"]}`),
            Build: func(a map[string]any) (string, string, any, error) {
                id, _ := a["id"].(string)
                if id == "" {
                    return "", "", nil, fmt.Errorf("id required")
                }
                return http.MethodGet, "/memories/" + id, nil, nil
            },
        },
        {
            Name:        "delete_memory",
            Description: "Delete a memory by track_id.",
            Schema:      json.RawMessage(`{"type":"object","properties":{"id":{"type":"string"}},"required":["id"]}`),
            Build: func(a map[string]any) (string, string, any, error) {
                id, _ := a["id"].(string)
                if id == "" {
                    return "", "", nil, fmt.Errorf("id required")
                }
                return http.MethodDelete, "/memories/" + id, nil, nil
            },
        },
        {
            Name:        "bulk_create_memories",
            Description: "Submit a batch of memories for async ingest.",
            Schema:      json.RawMessage(`{"type":"object","properties":{"items":{"type":"array","items":{"type":"object","properties":{"text":{"type":"string"},"metadata":{"type":"object"}},"required":["text"]}}},"required":["items"]}`),
            Build: func(a map[string]any) (string, string, any, error) {
                return http.MethodPost, "/memories:bulk", a, nil
            },
        },
        {
            Name:        "get_ingest_job",
            Description: "Poll the status of a bulk ingest job.",
            Schema:      json.RawMessage(`{"type":"object","properties":{"id":{"type":"string"}},"required":["id"]}`),
            Build: func(a map[string]any) (string, string, any, error) {
                id, _ := a["id"].(string)
                return http.MethodGet, "/ingest-jobs/" + id, nil, nil
            },
        },
        {
            Name:        "query",
            Description: "Retrieve memory references for the given query.",
            Schema:      json.RawMessage(`{"type":"object","properties":{"mode":{"type":"string","enum":["local","global","hybrid","naive","mix","bypass"]},"text":{"type":"string"},"top_k":{"type":"integer"}},"required":["mode","text"]}`),
            Build: func(a map[string]any) (string, string, any, error) {
                return http.MethodPost, "/queries", a, nil
            },
        },
        {
            Name:        "describe",
            Description: "Generative answer plus source paths for the given query.",
            Schema:      json.RawMessage(`{"type":"object","properties":{"mode":{"type":"string"},"text":{"type":"string"},"top_k":{"type":"integer"}},"required":["mode","text"]}`),
            Build: func(a map[string]any) (string, string, any, error) {
                return http.MethodPost, "/queries:describe", a, nil
            },
        },
        {
            Name:        "get_entity",
            Description: "Retrieve an entity by name.",
            Schema:      json.RawMessage(`{"type":"object","properties":{"id":{"type":"string"}},"required":["id"]}`),
            Build: func(a map[string]any) (string, string, any, error) {
                id, _ := a["id"].(string)
                return http.MethodGet, "/entities/" + id, nil, nil
            },
        },
        {
            Name:        "search_entities",
            Description: "Search entities by query string.",
            Schema:      json.RawMessage(`{"type":"object","properties":{"q":{"type":"string"}}}`),
            Build: func(a map[string]any) (string, string, any, error) {
                q, _ := a["q"].(string)
                path := "/entities"
                if q != "" {
                    path += "?q=" + q
                }
                return http.MethodGet, path, nil, nil
            },
        },
        {
            Name:        "patch_entity",
            Description: "Apply a partial update to an entity.",
            Schema:      json.RawMessage(`{"type":"object","properties":{"id":{"type":"string"},"patch":{"type":"object"}},"required":["id","patch"]}`),
            Build: func(a map[string]any) (string, string, any, error) {
                id, _ := a["id"].(string)
                patch := a["patch"]
                return http.MethodPatch, "/entities/" + id, patch, nil
            },
        },
        {
            Name:        "list_edges",
            Description: "List all edges in the knowledge graph.",
            Schema:      json.RawMessage(`{"type":"object","properties":{}}`),
            Build: func(a map[string]any) (string, string, any, error) {
                return http.MethodGet, "/edges", nil, nil
            },
        },
        {
            Name:        "create_edge",
            Description: "Create a new edge between two existing entities.",
            Schema:      json.RawMessage(`{"type":"object","properties":{"from_entity":{"type":"string"},"to_entity":{"type":"string"},"relation":{"type":"string"},"properties":{"type":"object"}},"required":["from_entity","to_entity","relation"]}`),
            Build: func(a map[string]any) (string, string, any, error) {
                return http.MethodPost, "/edges", a, nil
            },
        },
        {
            Name:        "delete_edge",
            Description: "Delete an edge by composite ID 'from||to'.",
            Schema:      json.RawMessage(`{"type":"object","properties":{"id":{"type":"string"}},"required":["id"]}`),
            Build: func(a map[string]any) (string, string, any, error) {
                id, _ := a["id"].(string)
                return http.MethodDelete, "/edges/" + id, nil, nil
            },
        },
    }
}

// Invoke runs the tool against the given client and returns the response body
// as raw JSON. Status >= 400 is surfaced as an error.
func Invoke(ctx context.Context, c *client.Client, t Tool, args map[string]any) ([]byte, error) {
    method, path, body, err := t.Build(args)
    if err != nil {
        return nil, err
    }
    resp, err := c.Do(ctx, method, path, body)
    if err != nil {
        return nil, err
    }
    defer resp.Body.Close()
    buf, _ := io.ReadAll(resp.Body)
    if resp.StatusCode >= 400 {
        return nil, fmt.Errorf("tatara: %s %s -> %d: %s", method, path, resp.StatusCode, strings.TrimSpace(string(buf)))
    }
    return buf, nil
}