1 files changed, 20 insertions, 1 deletions
diff --git a/packages/cli/src/ui/contexts/SessionContext.tsx b/packages/cli/src/ui/contexts/SessionContext.tsx
index 320df324..942af8b5 100644
--- a/packages/cli/src/ui/contexts/SessionContext.tsx
+++ b/packages/cli/src/ui/contexts/SessionContext.tsx
@@ -6,6 +6,7 @@
 
 import React, {
   createContext,
+  useCallback,
   useContext,
   useState,
   useMemo,
@@ -26,6 +27,7 @@ export interface SessionStatsState {
   sessionStartTime: Date;
   metrics: SessionMetrics;
   lastPromptTokenCount: number;
+  promptCount: number;
 }
 
 export interface ComputedSessionStats {
@@ -46,6 +48,8 @@ export interface ComputedSessionStats {
 // and the functions to update it.
 interface SessionStatsContextValue {
   stats: SessionStatsState;
+  startNewPrompt: () => void;
+  getPromptCount: () => number;
 }
 
 // --- Context Definition ---
@@ -63,6 +67,7 @@ export const SessionStatsProvider: React.FC<{ children: React.ReactNode }> = ({
     sessionStartTime: new Date(),
     metrics: uiTelemetryService.getMetrics(),
     lastPromptTokenCount: 0,
+    promptCount: 0,
   });
 
   useEffect(() => {
@@ -92,11 +97,25 @@ export const SessionStatsProvider: React.FC<{ children: React.ReactNode }> = ({
     };
   }, []);
 
+  const startNewPrompt = useCallback(() => {
+    setStats((prevState) => ({
+      ...prevState,
+      promptCount: prevState.promptCount + 1,
+    }));
+  }, []);
+
+  const getPromptCount = useCallback(
+    () => stats.promptCount,
+    [stats.promptCount],
+  );
+
   const value = useMemo(
     () => ({
       stats,
+      startNewPrompt,
+      getPromptCount,
     }),
-    [stats],
+    [stats, startNewPrompt, getPromptCount],
   );
 
   return (