Merge pull request #517 from underctrl-io/shared-store

feat: add shared store within environment

Author: twlite

apps/test-bot/src/app/commands/(general)/+ping.middleware.ts

@@ -10,16 +10,18 @@ export function beforeExecute(ctx: MiddlewareContext) {
   //   `Command-scoped middleware: ${ctx.commandName} will be executed!`,
   // );
 
-  Logger.info(`Command-scoped middleware: ${ctx.commandName} will be stopped!`);
-  Logger.info(
-    'None of the other beforeExecute middlewares are supposed to be executed',
-  );
+  ctx.store.set('foo', 'bar');
+
+  // Logger.info(`Command-scoped middleware: ${ctx.commandName} will be stopped!`);
+  // Logger.info(
+  //   'None of the other beforeExecute middlewares are supposed to be executed',
+  // );
 
-  after(() => {
-    Logger.info(`after() has been called in command-scoped middleware: ping`);
-  });
+  // after(() => {
+  //   Logger.info(`after() has been called in command-scoped middleware: ping`);
+  // });
 
-  stopMiddlewares();
+  // stopMiddlewares();
 }
 
 export function afterExecute(ctx: MiddlewareContext) {

apps/test-bot/src/app/commands/(general)/ping.ts

@@ -76,14 +76,15 @@ export async function autocomplete({
   interaction.respond(filtered);
 }
 
-export async function message({ message }: MessageCommandContext) {
-  message.reply('Pong!');
+export async function message(ctx: MessageCommandContext) {
+  Logger.debug`Store data ${ctx.store.get('foo')} | ${ctx.store}`;
+  ctx.message.reply('Pong!');
 }
 
-export async function chatInput({
-  interaction,
-  client,
-}: ChatInputCommandContext) {
+export async function chatInput(ctx: ChatInputCommandContext) {
+  const { interaction } = ctx;
+  Logger.debug`Store data ${ctx.store.get('foo')} | ${ctx.store}`;
+
   if (!interaction.channel) return;
 
   const button = new ButtonKit()

packages/ai/src/context.ts

@@ -1,5 +1,5 @@
 import type { CommandKit } from 'commandkit';
-import { Client, Message } from 'discord.js';
+import { Client, Collection, Message } from 'discord.js';
 
 /**
  * Options for the AI context.
@@ -48,7 +48,7 @@ export class AiContext<
   /**
    * A key-value store to hold additional data.
    */
-  public store = new Map<string, any>();
+  public store = new Collection<any, any>();
 
   /**
    * Creates a new instance of AiContext.

packages/commandkit/hooks.cjs

@@ -1,5 +1,6 @@
 const {
   useEnvironment,
+  useStore,
   getContext,
   getCommandKit,
 } = require('./dist/context/async-context.js');
@@ -51,6 +52,7 @@ function useEvent() {
 
 module.exports = {
   useAnyEnvironment,
+  useStore,
   useClient,
   useCommandKit,
   useCommand,

packages/commandkit/src/app/commands/Context.ts

@@ -3,6 +3,7 @@ import {
   Awaitable,
   ChatInputCommandInteraction,
   Client,
+  Collection,
   Guild,
   Interaction,
   Locale,
@@ -64,7 +65,7 @@ export interface ContextParameters<
   message: T extends 'message' ? Message : never;
   forwarded?: boolean;
   messageCommandParser?: T extends 'message' ? MessageCommandParser : never;
-  store?: Map<string, any>;
+  store?: Collection<string, any>;
   customArgs?: Args;
 }
 
@@ -244,7 +245,7 @@ export class Context<
    * @private
    * @internal
    */
-  #store: Map<string, any>;
+  #store: Collection<any, any>;
 
   /**
    * @private
@@ -265,7 +266,7 @@ export class Context<
     this.interaction = config.interaction;
     this.message = config.message;
     this.client = commandkit.client;
-    this.#store = config.store ?? new Map();
+    this.#store = config.environment?.store ?? config.store ?? new Collection();
     this.command = config.command;
 
     if (config.interaction) {
@@ -292,7 +293,7 @@ export class Context<
    * This store is shared across all contexts in the same command execution, including the cloned contexts and middleware contexts.
    */
   public get store() {
-    return this.#store;
+    return this.config.environment?.store ?? this.#store;
   }
 
   /**

packages/commandkit/src/context/async-context.ts

@@ -2,6 +2,7 @@ import { AsyncLocalStorage } from 'node:async_hooks';
 import { CommandKitEnvironment } from './environment';
 import { CommandKit } from '../commandkit';
 import { isCommandKitError } from '../utils/error-codes';
+import { Collection } from 'discord.js';
 
 const kCommandWorker = Symbol('commandkitCommandWorker');
 const context = new AsyncLocalStorage<CommandKitEnvironment>();
@@ -132,3 +133,10 @@ export function useEnvironment(): CommandKitEnvironment {
 
   return ctx;
 }
+
+/**
+ * Use the shared data store in the current environment. Throws an error if no context is found.
+ */
+export function useStore(): Collection<any, any> {
+  return useEnvironment().store;
+}

packages/commandkit/src/context/environment.ts

@@ -1,7 +1,8 @@
 import { randomUUID } from 'node:crypto';
 import { CommandKit } from '../commandkit';
-import { GenericFunction, getContext } from './async-context';
+import { GenericFunction, getContext, useEnvironment } from './async-context';
 import type { Context } from '../app';
+import { Collection } from 'discord.js';
 
 /**
  * Represents the internal data structure for the CommandKit environment.
@@ -45,6 +46,10 @@ export interface CommandKitEnvironmentInternalData {
    * This can be used to access request-specific data or application state.
    */
   context: Context | null;
+  /**
+   * Shared collection instance for arbitrary data storage
+   */
+  store: Collection<any, any>;
 }
 
 /**
@@ -60,6 +65,7 @@ export class CommandKitEnvironment {
     markStart: 0,
     markEnd: 0,
     context: null,
+    store: new Collection(),
   };
 
   /**
@@ -134,6 +140,13 @@ export class CommandKitEnvironment {
     return this.#data.variables;
   }
 
+  /**
+   * The shared store for this environment
+   */
+  public get store(): Collection<any, any> {
+    return this.#data.store;
+  }
+
   /**
    * Register a deferred function.
    * @param fn - The deferred function to register.

packages/tasks/src/context.ts

@@ -1,4 +1,4 @@
-import type { CommandKit, Client } from 'commandkit';
+import { type CommandKit, type Client, Collection } from 'commandkit';
 import { Task } from './task';
 
 /**
@@ -70,7 +70,7 @@ export class TaskContext<T extends Record<string, any> = Record<string, any>> {
    * });
    * ```
    */
-  public readonly store = new Map<string, any>();
+  public readonly store = new Collection<any, any>();
 
   /**
    * Creates a new task execution context.