Feature: Managing Agent Usage counter from backend (#293)

* feat: sending new header of last user message id on /v1/chat & storing counter in state.chat & resetting prevent_send on restored chat from cache

* wip: cut off enableSend

* fix: fixed types and tests

* fix: sending last user message id within body of request, not as header

* rework

* wip: receiving agent usage counter from backend, using those values in application

TODO: remove unused code from useAgentUsage() hook

* feat: refactored useAgentUsage hook, created middleware for chatResponse action dispatching, storing agent usage metadata in separate slice, not in chat one, receiving data on v1/login, asserting to state

* fix: created separate action for initial set of data, removed outdated comments

* fix: disabling agentic models for free users in other modes, not agent

* fix: adjusted UserSurvey test & made proper type guarding for User type

* chore: fixed npm run types

* fix: adjusted mocks for pro/non-pro users

---------

Co-authored-by: Kirill Starkov <[email protected]>

src/__fixtures__/msw.ts

@@ -66,6 +66,8 @@ export const goodUser: HttpHandler = http.get(
       inference: "PRO",
       metering_balance: -100000,
       questionnaire: {},
+      refact_agent_max_request_num: 20,
+      refact_agent_request_available: null,
     });
   },
 );
@@ -80,6 +82,8 @@ export const nonProUser: HttpHandler = http.get(
       inference: "FREE",
       metering_balance: -100000,
       questionnaire: {},
+      refact_agent_max_request_num: 20,
+      refact_agent_request_available: 5,
     });
   },
 );

src/__tests__/UserSurvey.test.tsx

@@ -25,9 +25,12 @@ const userMock = http.get(
       inference: "PRO",
       metering_balance: -100000,
       questionnaire: false,
+      refact_agent_max_request_num: 20,
+      refact_agent_request_available: null,
     });
   },
-  { once: true },
+  // TODO: if once if true, it still runs twice without refact_agent_max_request_num & refact_agent_request_available
+  // { once: true },
 );
 
 const questionnaireMock = http.get(

src/app/middleware.ts

@@ -10,6 +10,7 @@ import {
   chatAskQuestionThunk,
   restoreChat,
   newIntegrationChat,
+  chatResponse,
 } from "../features/Chat/Thread";
 import { statisticsApi } from "../services/refact/statistics";
 import { integrationsApi } from "../services/refact/integrations";
@@ -28,6 +29,14 @@ import { nextTip } from "../features/TipOfTheDay";
 import { telemetryApi } from "../services/refact/telemetry";
 import { CONFIG_PATH_URL, FULL_PATH_URL } from "../services/refact/consts";
 import { resetConfirmationInteractedState } from "../features/ToolConfirmation/confirmationSlice";
+import {
+  getAgentUsageCounter,
+  getMaxFreeAgentUsage,
+} from "../features/Chat/Thread/utils";
+import {
+  updateAgentUsage,
+  updateMaxAgentUsageAmount,
+} from "../features/AgentUsage/agentUsageSlice";
 
 export const listenerMiddleware = createListenerMiddleware();
 const startListening = listenerMiddleware.startListening.withTypes<
@@ -72,6 +81,31 @@ startListening({
   },
 });
 
+type ChatResponseAction = ReturnType<typeof chatResponse>;
+
+startListening({
+  matcher: isAnyOf((d: unknown): d is ChatResponseAction =>
+    chatResponse.match(d),
+  ),
+  effect: (action: ChatResponseAction, listenerApi) => {
+    const dispatch = listenerApi.dispatch;
+    // saving to store agent_usage counter from the backend, only one chunk has this field.
+    const { payload } = action;
+
+    if ("refact_agent_request_available" in payload) {
+      const agentUsageCounter = getAgentUsageCounter(payload);
+
+      dispatch(updateAgentUsage(agentUsageCounter ?? null));
+      // localStorage.setItem("agent_usage", agentUsageCounter?.toString() ?? "");
+    }
+
+    if ("refact_agent_max_request_num" in payload) {
+      const maxFreeAgentUsage = getMaxFreeAgentUsage(payload);
+      dispatch(updateMaxAgentUsageAmount(maxFreeAgentUsage));
+    }
+  },
+});
+
 startListening({
   // TODO: figure out why this breaks the tests when it's not a function :/
   matcher: isAnyOf(isRejected),

src/app/store.ts

@@ -54,10 +54,20 @@ const tipOfTheDayPersistConfig = {
   stateReconciler: mergeInitialState,
 };
 
+const agentUsagePersistConfig = {
+  key: "agentUsage",
+  storage: storage(),
+  stateReconciler: mergeInitialState,
+};
+
 const persistedTipOfTheDayReducer = persistReducer<
   ReturnType<typeof tipOfTheDaySlice.reducer>
 >(tipOfTheDayPersistConfig, tipOfTheDaySlice.reducer);
 
+const persistedAgentUsageReducer = persistReducer<
+  ReturnType<typeof agentUsageSlice.reducer>
+>(agentUsagePersistConfig, agentUsageSlice.reducer);
+
 // https://redux-toolkit.js.org/api/combineSlices
 // `combineSlices` automatically combines the reducers using
 // their `reducerPath`s, therefore we no longer need to call `combineReducers`.
@@ -67,6 +77,7 @@ const rootReducer = combineSlices(
     tour: tourReducer,
     // tipOfTheDay: persistedTipOfTheDayReducer,
     [tipOfTheDaySlice.reducerPath]: persistedTipOfTheDayReducer,
+    [agentUsageSlice.reducerPath]: persistedAgentUsageReducer,
     config: configReducer,
     active_file: activeFileReducer,
     selected_snippet: selectedSnippetReducer,
@@ -93,7 +104,6 @@ const rootReducer = combineSlices(
   attachedImagesSlice,
   userSurveySlice,
   integrationsSlice,
-  agentUsageSlice,
 );
 
 const rootPersistConfig = {

src/features/AgentUsage/AgentUsage.stories.tsx

@@ -6,18 +6,18 @@ import { Theme } from "../../components/Theme";
 import { AgentUsage } from "./AgentUsage";
 import { nonProUser } from "../../__fixtures__/msw";
 
-const items = Array.from({ length: 100 }).map(() => ({
-  user: "[email protected]",
-  time: Date.now(),
-}));
-
 const Template: React.FC = () => {
   const store = setUpStore({
     tour: {
       type: "finished",
     },
     agentUsage: {
-      items,
+      agent_usage: 5,
+      agent_max_usage_amount: 20,
+      _persist: {
+        rehydrated: true,
+        version: 1,
+      },
     },
     config: {
       apiKey: "foo",

src/features/AgentUsage/AgentUsage.tsx

@@ -1,34 +1,33 @@
 import React, { useMemo } from "react";
 
-import { useAgentUsage, useGetUser } from "../../hooks";
+import { useAgentUsage, useAppSelector, useGetUser } from "../../hooks";
 import { Flex, Card, Text } from "@radix-ui/themes";
 import { LinkButton } from "../../components/Buttons";
 import styles from "./AgentUsage.module.css";
+import { selectAgentUsage } from "./agentUsageSlice";
 
 export const AgentUsage: React.FC = () => {
   const userRequest = useGetUser();
+  const agentUsageAmount = useAppSelector(selectAgentUsage);
 
-  const { usersUsage, shouldShow, MAX_FREE_USAGE, startPollingForUser, plan } =
+  const { shouldShow, maxAgentUsageAmount, startPollingForUser, plan } =
     useAgentUsage();
 
   const usageMessage = useMemo(() => {
-    if (usersUsage >= MAX_FREE_USAGE) {
-      return `You have reached your usage limit of ${MAX_FREE_USAGE} messages a day.
+    if (agentUsageAmount === null) return null;
+    if (agentUsageAmount === 0) {
+      return `You have reached your usage limit of ${maxAgentUsageAmount} messages a day.
           You can use agent again tomorrow, or upgrade to PRO.`;
     }
 
-    if (usersUsage >= MAX_FREE_USAGE - 5) {
-      return `You have left only ${
-        MAX_FREE_USAGE - usersUsage
-      } messages left today. To remove
+    if (agentUsageAmount <= 5) {
+      return `You have left only ${agentUsageAmount} messages left today. To remove
           the limit upgrade to PRO.`;
     }
 
-    return `You have ${
-      MAX_FREE_USAGE - usersUsage
-    } agent messages left on our ${plan}
+    return `You have ${agentUsageAmount} agent messages left on our ${plan}
         plan.`;
-  }, [MAX_FREE_USAGE, plan, usersUsage]);
+  }, [maxAgentUsageAmount, plan, agentUsageAmount]);
 
   if (!userRequest.data) return null;
   if (!shouldShow) return null;

src/features/AgentUsage/agentUsageSlice.ts

@@ -1,34 +1,45 @@
 import { createSlice, PayloadAction } from "@reduxjs/toolkit";
 
-export type AgentUsageItem = {
-  time: number;
-  user: string;
+export type AgentUsageMeta = {
+  agent_usage: null | number; // null if plan is PRO or ROBOT
+  agent_max_usage_amount: number; // maximum amount of agent usage allowed per UTC day for users with FREE plan
 };
 
-const initialState: { items: AgentUsageItem[] } = { items: [] };
-
-const oneDay = 24 * 60 * 60 * 1000;
+const initialState: AgentUsageMeta = {
+  agent_usage: null,
+  agent_max_usage_amount: 20,
+};
 
 export const agentUsageSlice = createSlice({
   name: "agentUsage",
   initialState,
   reducers: {
-    addAgentUsageItem: (state, action: PayloadAction<{ user: string }>) => {
-      const now = Date.now();
-      const todaysItems = state.items.filter(
-        (item) => item.time + oneDay > now,
-      );
-      const item = { time: now, user: action.payload.user };
-      state.items = [...todaysItems, item];
+    updateAgentUsage: (
+      state,
+      action: PayloadAction<AgentUsageMeta["agent_usage"]>,
+    ) => {
+      state.agent_usage = action.payload;
+    },
+    updateMaxAgentUsageAmount: (state, action: PayloadAction<number>) => {
+      state.agent_max_usage_amount = action.payload;
+    },
+    setInitialAgentUsage: (state, action: PayloadAction<AgentUsageMeta>) => {
+      const { agent_max_usage_amount, agent_usage } = action.payload;
+      state.agent_usage = agent_usage;
+      state.agent_max_usage_amount = agent_max_usage_amount;
     },
   },
 
   selectors: {
-    selectAgentUsageItems: (state) => {
-      return state.items;
-    },
+    selectAgentUsage: (state) => state.agent_usage,
+    selectMaxAgentUsageAmount: (state) => state.agent_max_usage_amount,
   },
 });
 
-export const { addAgentUsageItem } = agentUsageSlice.actions;
-export const { selectAgentUsageItems } = agentUsageSlice.selectors;
+export const {
+  updateAgentUsage,
+  updateMaxAgentUsageAmount,
+  setInitialAgentUsage,
+} = agentUsageSlice.actions;
+export const { selectAgentUsage, selectMaxAgentUsageAmount } =
+  agentUsageSlice.selectors;

src/features/Chat/Thread/actions.ts

@@ -6,6 +6,7 @@ import {
   type ToolUse,
   IntegrationMeta,
   LspChatMode,
+  PayloadWithChatAndMessageId,
 } from "./types";
 import {
   isAssistantDelta,
@@ -51,6 +52,10 @@ export const chatAskedQuestion = createAction<PayloadWithId>(
   "chatThread/askQuestion",
 );
 
+export const setLastUserMessageId = createAction<PayloadWithChatAndMessageId>(
+  "chatThread/setLastUserMessageId",
+);
+
 export const backUpMessages = createAction<
   PayloadWithId & {
     messages: ChatThread["messages"];
@@ -279,8 +284,11 @@ export const chatAskQuestionThunk = createAppAsyncThunk<
 
     const messagesForLsp = formatMessagesForLsp(messages);
     const realMode = mode ?? thread?.mode;
+    const maybeLastUserMessageId = thread?.last_user_message_id;
+
     return sendChat({
       messages: messagesForLsp,
+      last_user_message_id: maybeLastUserMessageId,
       model: state.chat.thread.model,
       tools,
       stream: true,

src/features/Chat/Thread/reducer.ts

@@ -30,6 +30,7 @@ import {
   setIsWaitingForResponse,
   setMaxNewTokens,
   setAutomaticPatch,
+  setLastUserMessageId,
 } from "./actions";
 import { formatChatResponse } from "./utils";
 import { DEFAULT_MAX_NEW_TOKENS } from "../../../services/refact";
@@ -44,6 +45,7 @@ const createChatThread = (
     messages: [],
     title: "",
     model: "",
+    last_user_message_id: "",
     tool_use,
     integration,
     mode,
@@ -164,6 +166,11 @@ export const chatReducer = createReducer(initialState, (builder) => {
     state.automatic_patch = action.payload;
   });
 
+  builder.addCase(setLastUserMessageId, (state, action) => {
+    if (state.thread.id !== action.payload.chatId) return state;
+    state.thread.last_user_message_id = action.payload.messageId;
+  });
+
   builder.addCase(chatAskedQuestion, (state, action) => {
     if (state.thread.id !== action.payload.id) return state;
     state.send_immediately = false;

src/features/Chat/Thread/types.ts

@@ -20,6 +20,7 @@ export type ChatThread = {
   isTitleGenerated?: boolean;
   integration?: IntegrationMeta | null;
   mode?: LspChatMode;
+  last_user_message_id?: string;
 };
 
 export type ToolUse = "quick" | "explore" | "agent";
@@ -39,6 +40,7 @@ export type Chat = {
 };
 
 export type PayloadWithId = { id: string };
+export type PayloadWithChatAndMessageId = { chatId: string; messageId: string };
 export type PayloadWithIdAndTitle = {
   title: string;
   isTitleGenerated: boolean;

src/features/Chat/Thread/utils.ts

@@ -100,7 +100,7 @@ export function mergeToolCalls(prev: ToolCall[], add: ToolCall[]): ToolCall[] {
   }, prev);
 }
 
-function lastIndexOf<T>(arr: T[], predicate: (a: T) => boolean): number {
+export function lastIndexOf<T>(arr: T[], predicate: (a: T) => boolean): number {
   let index = -1;
   for (let i = arr.length - 1; i >= 0; i--) {
     if (predicate(arr[i])) {
@@ -128,6 +128,20 @@ function replaceLastUserMessage(
   return result.concat([userMessage]);
 }
 
+export function getAgentUsageCounter(response: ChatResponse): number | null {
+  if (isChatResponseChoice(response)) {
+    return response.refact_agent_request_available;
+  }
+  return null;
+}
+
+export function getMaxFreeAgentUsage(response: ChatResponse): number {
+  if (isChatResponseChoice(response)) {
+    return response.refact_agent_max_request_num;
+  }
+  return 0;
+}
+
 export function formatChatResponse(
   messages: ChatMessages,
   response: ChatResponse,