diff --git a/cmd/server/main.go b/cmd/server/main.go index 9653afd..15e6ea1 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -23,6 +23,7 @@ import ( "github.com/tiennm99/miti99bot/internal/modules/loldle" "github.com/tiennm99/miti99bot/internal/modules/lolschedule" "github.com/tiennm99/miti99bot/internal/modules/misc" + "github.com/tiennm99/miti99bot/internal/modules/stats" "github.com/tiennm99/miti99bot/internal/modules/trading" "github.com/tiennm99/miti99bot/internal/modules/twentyq" "github.com/tiennm99/miti99bot/internal/modules/util" @@ -49,6 +50,7 @@ func factories() map[string]modules.Factory { "lolschedule": lolschedule.New, "twentyq": twentyq.New, "trading": trading.New, + "stats": stats.New, } } diff --git a/internal/modules/dispatcher.go b/internal/modules/dispatcher.go index 6df03d9..9be4e0b 100644 --- a/internal/modules/dispatcher.go +++ b/internal/modules/dispatcher.go @@ -3,6 +3,7 @@ package modules import ( "context" "strings" + "time" "github.com/go-telegram/bot" "github.com/go-telegram/bot/models" @@ -66,6 +67,11 @@ func Install(b *bot.Bot, reg *Registry, auth Auth) { return // silent — do not leak existence of gated commands } metrics.IncCommand(cmdCopy.Name) + go func() { + hookCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + reg.RunCommandHooks(hookCtx, cmdCopy.Name) + }() if err := cmdCopy.Handler(ctx, b, update); err != nil { metrics.IncError("handler-error") log.Error("command failed", "command", cmdCopy.Name, "err", err) diff --git a/internal/modules/module.go b/internal/modules/module.go index 0fcb2bd..a647fa6 100644 --- a/internal/modules/module.go +++ b/internal/modules/module.go @@ -55,9 +55,10 @@ type Cron struct { // Module.Name is overridden by the registry to its catalog key; factories may // leave it blank. type Module struct { - Name string - Commands []Command - Crons []Cron + Name string + Commands []Command + Crons []Cron + CommandHook func(ctx context.Context, name string) // optional; called by dispatcher after each authorized command invocation } // Deps is the dependency bundle a Factory receives. diff --git a/internal/modules/registry.go b/internal/modules/registry.go index 88a627a..aa47768 100644 --- a/internal/modules/registry.go +++ b/internal/modules/registry.go @@ -1,6 +1,7 @@ package modules import ( + "context" "fmt" "regexp" "sort" @@ -28,13 +29,14 @@ var moduleNameRe = regexp.MustCompile(`^[a-z0-9_-]{1,32}$`) // dispatchers and handlers capture *Registry by pointer and assume the maps // are stable. A future hot-reload feature would need an explicit mutation API. type Registry struct { - Modules []Module // in MODULES-env order - AllCommands map[string]Command // name → Command, deduped across modules - publicCmds map[string]Command - protected map[string]Command - private map[string]Command - crons map[string]Cron // name → Cron, unique across modules - cronDeps map[string]Deps // cron name → owning module's prefixed Deps + Modules []Module // in MODULES-env order + AllCommands map[string]Command // name → Command, deduped across modules + publicCmds map[string]Command + protected map[string]Command + private map[string]Command + crons map[string]Cron // name → Cron, unique across modules + cronDeps map[string]Deps // cron name → owning module's prefixed Deps + commandHooks []func(ctx context.Context, name string) } // PublicCommands returns commands tagged VisibilityPublic, sorted by name. @@ -59,6 +61,15 @@ func (r *Registry) CronDeps(name string) (Deps, bool) { return d, ok } +// RunCommandHooks calls every CommandHook registered by loaded modules in +// order. Errors are not returned — hooks are best-effort (e.g., stats +// counters) and must not fail the command handler. +func (r *Registry) RunCommandHooks(ctx context.Context, name string) { + for _, h := range r.commandHooks { + h(ctx, name) + } +} + // Crons returns all loaded crons, sorted by name. Allocates a fresh slice on // every call — fine for startup-time logging, not for hot paths. func (r *Registry) Crons() []Cron { @@ -136,6 +147,9 @@ func Build(enabled []string, factories map[string]Factory, kv storage.KVProvider return nil, fmt.Errorf("module %q: factory returned mismatched Name=%q", name, mod.Name) } mod.Name = name + if mod.CommandHook != nil { + reg.commandHooks = append(reg.commandHooks, mod.CommandHook) + } for _, cmd := range mod.Commands { if err := validateCommand(cmd); err != nil { diff --git a/internal/modules/stats/stats.go b/internal/modules/stats/stats.go new file mode 100644 index 0000000..9cd6c23 --- /dev/null +++ b/internal/modules/stats/stats.go @@ -0,0 +1,116 @@ +// Package stats tracks per-command invocation counts persistently in KV and +// exposes /stats to display them sorted by popularity. +package stats + +import ( + "context" + "errors" + "fmt" + "sort" + "strings" + + "github.com/go-telegram/bot" + "github.com/go-telegram/bot/models" + + "github.com/tiennm99/miti99bot/internal/log" + "github.com/tiennm99/miti99bot/internal/modules" + "github.com/tiennm99/miti99bot/internal/modules/util/chathelper" + "github.com/tiennm99/miti99bot/internal/storage" +) + +const countPrefix = "count:" + +type countEntry struct { + N int64 `json:"n"` +} + +type counter struct { + kv storage.KVStore +} + +func countKey(name string) string { return countPrefix + name } + +// Inc increments the persistent invocation count for the named command. +// Errors are logged and swallowed — stats are best-effort. +func (c *counter) Inc(ctx context.Context, name string) { + key := countKey(name) + var entry countEntry + if err := c.kv.GetJSON(ctx, key, &entry); err != nil && !errors.Is(err, storage.ErrNotFound) { + log.Error("stats: kv get failed", "key", key, "err", err) + return + } + entry.N++ + if err := c.kv.PutJSON(ctx, key, entry); err != nil { + log.Error("stats: kv put failed", "key", key, "err", err) + } +} + +// New is the module Factory. Registers a CommandHook that persists counts and +// a /stats command that displays them. +func New(deps modules.Deps) modules.Module { + c := &counter{kv: deps.KV} + return modules.Module{ + CommandHook: c.Inc, + Commands: []modules.Command{ + statsCommand(c), + }, + } +} + +func statsCommand(c *counter) modules.Command { + return modules.Command{ + Name: "stats", + Visibility: modules.VisibilityPublic, + Description: "Show command usage statistics", + Handler: func(ctx context.Context, b *bot.Bot, update *models.Update) error { + if update.Message == nil { + return nil + } + keys, err := c.kv.List(ctx, countPrefix) + if err != nil { + log.Error("stats: kv list failed", "err", err) + return chathelper.Reply(ctx, b, update.Message, "Could not load stats. Try again later.") + } + if len(keys) == 0 { + return chathelper.Reply(ctx, b, update.Message, "No command stats yet.") + } + + type row struct { + name string + n int64 + } + rows := make([]row, 0, len(keys)) + for _, k := range keys { + name := strings.TrimPrefix(k, countPrefix) + var entry countEntry + if err := c.kv.GetJSON(ctx, k, &entry); err != nil { + log.Error("stats: kv get failed during render", "key", k, "err", err) + continue + } + rows = append(rows, row{name: name, n: entry.N}) + } + sort.Slice(rows, func(i, j int) bool { + if rows[i].n != rows[j].n { + return rows[i].n > rows[j].n + } + return rows[i].name < rows[j].name + }) + + var sb strings.Builder + sb.WriteString("Command usage:\n") + for _, r := range rows { + fmt.Fprintf(&sb, "/%s: %d\n", r.name, r.n) + } + const telegramMaxLen = 4000 // leave margin below Telegram's 4096-byte hard limit + output := strings.TrimSuffix(sb.String(), "\n") + if len(output) > telegramMaxLen { + cutoff := strings.LastIndexByte(output[:telegramMaxLen], '\n') + if cutoff <= 0 { + cutoff = telegramMaxLen + } + output = output[:cutoff] + "\n…(truncated)" + } + return chathelper.Reply(ctx, b, update.Message, output) + }, + } +} diff --git a/internal/modules/stats/stats_test.go b/internal/modules/stats/stats_test.go new file mode 100644 index 0000000..d7f5e80 --- /dev/null +++ b/internal/modules/stats/stats_test.go @@ -0,0 +1,151 @@ +package stats + +import ( + "context" + "strings" + "testing" + + "github.com/tiennm99/miti99bot/internal/modules" + "github.com/tiennm99/miti99bot/internal/storage" + "github.com/tiennm99/miti99bot/internal/testutil" +) + +func TestNew_RegistersExpectedCommands(t *testing.T) { + deps := modules.Deps{KV: storage.NewMemoryKVStore()} + mod := New(deps) + + if len(mod.Commands) != 1 { + t.Fatalf("commands count = %d, want 1", len(mod.Commands)) + } + cmd := mod.Commands[0] + if cmd.Name != "stats" { + t.Errorf("command name = %q, want %q", cmd.Name, "stats") + } + if cmd.Visibility != modules.VisibilityPublic { + t.Errorf("command visibility = %d, want Public", cmd.Visibility) + } + if cmd.Handler == nil { + t.Error("command handler is nil") + } + if mod.CommandHook == nil { + t.Error("CommandHook is nil") + } +} + +func TestInc_PersistsCountInKV(t *testing.T) { + ctx := context.Background() + kv := storage.NewMemoryKVStore() + c := &counter{kv: kv} + + c.Inc(ctx, "ping") + c.Inc(ctx, "ping") + c.Inc(ctx, "wordle") + + var entry countEntry + if err := kv.GetJSON(ctx, countKey("ping"), &entry); err != nil { + t.Fatalf("GetJSON ping: %v", err) + } + if entry.N != 2 { + t.Errorf("ping count = %d, want 2", entry.N) + } + + entry = countEntry{} + if err := kv.GetJSON(ctx, countKey("wordle"), &entry); err != nil { + t.Fatalf("GetJSON wordle: %v", err) + } + if entry.N != 1 { + t.Errorf("wordle count = %d, want 1", entry.N) + } +} + +func installStats(t *testing.T) (*testutil.RecordingBot, *counter) { + t.Helper() + rb := testutil.NewRecordingBot(t) + kv := storage.NewMemoryKVStore() + c := &counter{kv: kv} + mod := modules.Module{ + Commands: []modules.Command{statsCommand(c)}, + } + reg := &modules.Registry{ + AllCommands: map[string]modules.Command{}, + } + for _, cmd := range mod.Commands { + reg.AllCommands[cmd.Name] = cmd + } + modules.Install(rb.Bot, reg, modules.Auth{}) + return rb, c +} + +func TestStats_NoDataRepliesEmpty(t *testing.T) { + rb, _ := installStats(t) + rb.Bot.ProcessUpdate(context.Background(), testutil.NewPrivateMessage(1, "/stats")) + + got := rb.LastSent().Text() + if got != "No command stats yet." { + t.Errorf("empty stats reply = %q, want 'No command stats yet.'", got) + } +} + +func TestStats_ShowsCountsSortedByPopularity(t *testing.T) { + ctx := context.Background() + rb, c := installStats(t) + + c.Inc(ctx, "ping") + c.Inc(ctx, "wordle") + c.Inc(ctx, "wordle") + c.Inc(ctx, "wordle") + c.Inc(ctx, "loldle") + c.Inc(ctx, "loldle") + + rb.Bot.ProcessUpdate(ctx, testutil.NewPrivateMessage(1, "/stats")) + got := rb.LastSent().Text() + + if !strings.HasPrefix(got, "Command usage:") { + t.Errorf("reply missing header: %q", got) + } + // Verify descending order: wordle(3) > loldle(2) > ping(1) + wordlePos := strings.Index(got, "/wordle:") + loLdlePos := strings.Index(got, "/loldle:") + pingPos := strings.Index(got, "/ping:") + if wordlePos < 0 || loLdlePos < 0 || pingPos < 0 { + t.Fatalf("reply missing expected commands: %q", got) + } + if !(wordlePos < loLdlePos && loLdlePos < pingPos) { + t.Errorf("commands not in descending count order: wordle=%d loldle=%d ping=%d in %q", + wordlePos, loLdlePos, pingPos, got) + } +} + +func TestCommandHook_FiredThroughModulesBuild(t *testing.T) { + ctx := context.Background() + provider := storage.NewMemoryProvider() + + reg, err := modules.Build( + []string{"stats"}, + map[string]modules.Factory{"stats": New}, + provider, + modules.BuildOptions{}, + ) + if err != nil { + t.Fatalf("Build: %v", err) + } + if len(reg.Modules) != 1 { + t.Fatalf("expected 1 module, got %d", len(reg.Modules)) + } + + rb := testutil.NewRecordingBot(t) + modules.Install(rb.Bot, reg, modules.Auth{}) + + // Dispatch /ping — not a registered command, so nothing replies. + // But RunCommandHooks should have fired and incremented count:ping. + reg.RunCommandHooks(ctx, "ping") + + statsKV := provider.For("stats") + var entry countEntry + if err := statsKV.GetJSON(ctx, countKey("ping"), &entry); err != nil { + t.Fatalf("expected count:ping in KV after hook: %v", err) + } + if entry.N != 1 { + t.Errorf("count:ping = %d, want 1", entry.N) + } +} diff --git a/template.yaml b/template.yaml index aed85d8..71c4fdd 100644 --- a/template.yaml +++ b/template.yaml @@ -14,7 +14,7 @@ Parameters: ModulesCSV: Type: String - Default: util,misc,wordle,loldle,lolschedule,twentyq,trading + Default: util,misc,wordle,loldle,lolschedule,twentyq,trading,stats Description: Comma-separated module names enabled at runtime (matches MODULES env). BotOwnerID: @@ -135,6 +135,8 @@ Resources: TELEGRAM_WEBHOOK_SECRET_PARAMETER_NAME: !Sub "/miti99bot/${StackEnv}/telegram-webhook-secret" GEMINI_API_KEY_PARAMETER_NAME: !Sub "/miti99bot/${StackEnv}/gemini-api-key" CRON_SHARED_SECRET_PARAMETER_NAME: !Sub "/miti99bot/${StackEnv}/cron-shared-secret" + # stats module uses non-atomic KV increment; concurrency > 1 would lose counts. + ReservedConcurrentExecutions: 1 FunctionUrlConfig: AuthType: NONE InvokeMode: BUFFERED