@pensar/apex 0.0.37-canary.0 → 0.0.39-canary.2f181ec5

package/build/index.js

@@ -27849,24 +27849,40 @@ function useAgent() {
 }
 function AgentProvider({ children }) {
   const [model, setModel] = import_react10.useState(AVAILABLE_MODELS[0]);
-  const [tokenCount, setTokenCount] = import_react10.useState(0);
+  const [tokenUsage, setTokenUsage] = import_react10.useState({
+    inputTokens: 0,
+    outputTokens: 0,
+    totalTokens: 0
+  });
+  const [hasExecuted, setHasExecuted] = import_react10.useState(false);
   const [thinking, setThinking] = import_react10.useState(false);
   const [isExecuting, setIsExecuting] = import_react10.useState(false);
-  const addTokens = (tokens) => {
-    setTokenCount((prev) => prev + tokens);
-  };
+  const addTokenUsage = import_react10.useCallback((input, output) => {
+    setHasExecuted(true);
+    setTokenUsage((prev) => ({
+      inputTokens: prev.inputTokens + input,
+      outputTokens: prev.outputTokens + output,
+      totalTokens: prev.totalTokens + input + output
+    }));
+  }, []);
+  const resetTokenUsage = import_react10.useCallback(() => {
+    setHasExecuted(false);
+    setTokenUsage({ inputTokens: 0, outputTokens: 0, totalTokens: 0 });
+  }, []);
+  const contextValue = import_react10.useMemo(() => ({
+    model,
+    setModel,
+    tokenUsage,
+    addTokenUsage,
+    resetTokenUsage,
+    hasExecuted,
+    thinking,
+    setThinking,
+    isExecuting,
+    setIsExecuting
+  }), [model, tokenUsage, hasExecuted, thinking, isExecuting, addTokenUsage, resetTokenUsage]);
   return /* @__PURE__ */ import_jsx_dev_runtime2.jsxDEV(AgentContext.Provider, {
-    value: {
-      model,
-      setModel,
-      tokenCount,
-      setTokenCount,
-      addTokens,
-      thinking,
-      setThinking,
-      isExecuting,
-      setIsExecuting
-    },
+    value: contextValue,
     children
   }, undefined, false, undefined, this);
 }
@@ -27940,7 +27956,7 @@ function Footer({
   showExitWarning = false
 }) {
   cwd = "~" + cwd.split(os2.homedir()).pop() || "";
-  const { model, tokenCount, thinking, isExecuting } = useAgent();
+  const { model, tokenUsage, hasExecuted, thinking, isExecuting } = useAgent();
   const hotkeys = isExecuting ? [{ key: "Ctrl+C", label: "Stop Execution" }] : [{ key: "Ctrl+C", label: "Clear/Exit" }];
   return /* @__PURE__ */ import_jsx_dev_runtime2.jsxDEV("box", {
     flexDirection: "row",
@@ -28006,27 +28022,21 @@ function Footer({
   }, undefined, true, undefined, this);
 }
 function AgentStatus() {
-  const { tokenCount, thinking, isExecuting } = useAgent();
+  const { tokenUsage, hasExecuted, thinking, isExecuting } = useAgent();
   return /* @__PURE__ */ import_jsx_dev_runtime2.jsxDEV("box", {
     flexDirection: "row",
     gap: 1,
     children: [
-      tokenCount > 0 && /* @__PURE__ */ import_jsx_dev_runtime2.jsxDEV(import_jsx_dev_runtime2.Fragment, {
+      hasExecuted && /* @__PURE__ */ import_jsx_dev_runtime2.jsxDEV(import_jsx_dev_runtime2.Fragment, {
         children: [
           /* @__PURE__ */ import_jsx_dev_runtime2.jsxDEV("box", {
             border: ["right"],
             borderColor: "green"
           }, undefined, false, undefined, this),
           /* @__PURE__ */ import_jsx_dev_runtime2.jsxDEV("text", {
-            fg: "gray",
-            children: [
-              "■ ",
-              /* @__PURE__ */ import_jsx_dev_runtime2.jsxDEV("span", {
-                fg: "white",
-                children: formatTokenCount(tokenCount)
-              }, undefined, false, undefined, this)
-            ]
-          }, undefined, true, undefined, this),
+            fg: "white",
+            children: `↓${formatTokenCount(tokenUsage.inputTokens)} ↑${formatTokenCount(tokenUsage.outputTokens)} Σ${formatTokenCount(tokenUsage.totalTokens)}`
+          }, undefined, false, undefined, this),
           /* @__PURE__ */ import_jsx_dev_runtime2.jsxDEV(ContextProgress, {
             width: 10
           }, undefined, false, undefined, this)
@@ -28048,10 +28058,11 @@ function AgentStatus() {
   }, undefined, true, undefined, this);
 }
 function ContextProgress({ width }) {
-  const { model, tokenCount, thinking } = useAgent();
-  const contextProgress = Number((tokenCount / (model.contextLength ?? 200000) * 100).toFixed(2));
-  if (!thinking)
+  const { model, tokenUsage, thinking } = useAgent();
+  if (!thinking || tokenUsage.totalTokens === 0)
     return null;
+  const contextLength = model.contextLength ?? 200000;
+  const contextProgress = Math.max(0, Math.min(100, Number((tokenUsage.totalTokens / contextLength * 100).toFixed(2))));
   return /* @__PURE__ */ import_jsx_dev_runtime2.jsxDEV(ProgressBar, {
     value: contextProgress,
     width
@@ -65017,11 +65028,38 @@ async function summarizeConversation(messages, opts, model) {
       content: `Summarize this conversation to pass to another agent. This was the system prompt: ${opts.system} `
     }
   ];
-  const { text: summary } = await generateText({
+  const { text: summary, usage: summaryUsage } = await generateText({
     model,
     system: `You are a helpful assistant that summarizes conversations to pass to another agent. Review the conversation and system prompt at the end provided by the user.`,
     messages: summarizedMessages
   });
+  if (opts.onStepFinish && summaryUsage) {
+    opts.onStepFinish({
+      text: "",
+      reasoning: undefined,
+      reasoningDetails: [],
+      files: [],
+      sources: [],
+      toolCalls: [],
+      toolResults: [],
+      finishReason: "stop",
+      usage: {
+        inputTokens: summaryUsage.inputTokens ?? 0,
+        outputTokens: summaryUsage.outputTokens ?? 0,
+        totalTokens: summaryUsage.totalTokens ?? 0
+      },
+      warnings: [],
+      request: {},
+      response: {
+        id: "summarization",
+        timestamp: new Date,
+        modelId: ""
+      },
+      providerMetadata: undefined,
+      stepType: "initial",
+      isContinued: false
+    });
+  }
   const originalLength = typeof opts.prompt === "string" ? opts.prompt.length : 0;
   const enhancedPrompt = originalLength > 1e5 ? `Context: The previous conversation contained very long content that was summarized.
 
@@ -65184,6 +65222,7 @@ function streamResponse(opts) {
   } = opts;
   const messagesContainer = { current: messages || [] };
   const providerModel = getProviderModel(model, authConfig);
+  let rateLimitRetryCount = 0;
   try {
     const response = streamText({
       model: providerModel,
@@ -65197,6 +65236,16 @@ function streamResponse(opts) {
         messagesContainer.current = opts2.messages;
         return;
       },
+      onError: async ({ error: error46 }) => {
+        if (error46.message.toLowerCase().includes("too many tokens") || error46.message.toLowerCase().includes("overloaded")) {
+          rateLimitRetryCount++;
+          await new Promise((resolve4) => setTimeout(resolve4, 1000 * rateLimitRetryCount));
+          if (rateLimitRetryCount < 20) {
+            return;
+          }
+        }
+        throw error46;
+      },
       onStepFinish,
       abortSignal,
       activeTools,
@@ -65215,7 +65264,7 @@ function streamResponse(opts) {
             throw new Error(`Tool ${toolCall.toolName} not found or has no schema`);
           }
           const jsonSchema2 = inputSchema({ toolName: toolCall.toolName });
-          const { object: repairedArgs } = await generateObject({
+          const { object: repairedArgs, usage: repairUsage } = await generateObject({
             model: providerModel,
             schema: tool2.inputSchema,
             prompt: [
@@ -65228,6 +65277,33 @@ function streamResponse(opts) {
             ].join(`
 `)
           });
+          if (onStepFinish && repairUsage) {
+            onStepFinish({
+              text: "",
+              reasoning: undefined,
+              reasoningDetails: [],
+              files: [],
+              sources: [],
+              toolCalls: [],
+              toolResults: [],
+              finishReason: "stop",
+              usage: {
+                inputTokens: repairUsage.inputTokens ?? 0,
+                outputTokens: repairUsage.outputTokens ?? 0,
+                totalTokens: repairUsage.totalTokens ?? 0
+              },
+              warnings: [],
+              request: {},
+              response: {
+                id: "tool-repair",
+                timestamp: new Date,
+                modelId: ""
+              },
+              providerMetadata: undefined,
+              stepType: "initial",
+              isContinued: false
+            });
+          }
           return { ...toolCall, input: JSON.stringify(repairedArgs) };
         } catch (repairError) {
           if (!silent) {
@@ -65254,9 +65330,9 @@ function streamResponse(opts) {
   }
 }
 async function generateObjectResponse(opts) {
-  const { model, schema, prompt, system, maxTokens, temperature, authConfig } = opts;
+  const { model, schema, prompt, system, maxTokens, temperature, authConfig, onTokenUsage } = opts;
   const providerModel = getProviderModel(model, authConfig);
-  const { object: object3 } = await generateObject({
+  const { object: object3, usage } = await generateObject({
     model: providerModel,
     schema,
     prompt,
@@ -65264,6 +65340,9 @@ async function generateObjectResponse(opts) {
     maxTokens,
     temperature
   });
+  if (onTokenUsage && usage) {
+    onTokenUsage(usage.inputTokens ?? 0, usage.outputTokens ?? 0);
+  }
   return object3;
 }
 // src/core/agent/pentestAgent/prompts.ts
@@ -68645,7 +68724,7 @@ Example workflow:
     execute: async (params) => recordTestResultCore(session, params)
   });
 }
-async function generateTestStrategy(params, model) {
+async function generateTestStrategy(params, model, onTokenUsage) {
   const prompt = `You are a penetration testing expert. Generate a concise testing strategy:
 
 Attack Type: ${params.knowledge.name}
@@ -68671,12 +68750,15 @@ Be tactical and specific.`;
       model: providerModel,
       prompt
     });
+    if (onTokenUsage && result.usage) {
+      onTokenUsage(result.usage.inputTokens ?? 0, result.usage.outputTokens ?? 0);
+    }
     return result.text;
   } catch (error46) {
     return params.knowledge.adaptiveStrategy;
   }
 }
-async function generatePayload(params, model) {
+async function generatePayload(params, model, onTokenUsage) {
   const prompt = `Generate ONE ${params.knowledge.name} payload for testing.
 
 Techniques:
@@ -68700,7 +68782,8 @@ Generate ONE specific payload. Return ONLY JSON:
     const result = await generateObjectResponse({
       model,
       schema: PayloadSchema,
-      prompt
+      prompt,
+      onTokenUsage
     });
     return result;
   } catch (error46) {
@@ -68713,7 +68796,7 @@ Generate ONE specific payload. Return ONLY JSON:
     technique: technique.name
   };
 }
-async function analyzeResponse(params, model) {
+async function analyzeResponse(params, model, onTokenUsage) {
   const prompt = `Analyze this security test response:
 
 Attack: ${params.knowledge.name}
@@ -68741,7 +68824,8 @@ Analyze: Is this vulnerable? Return ONLY JSON:
     const result = await generateObjectResponse({
       model,
       schema: AnalysisSchema,
-      prompt
+      prompt,
+      onTokenUsage
     });
     return result;
   } catch (error46) {
@@ -68760,7 +68844,7 @@ Analyze: Is this vulnerable? Return ONLY JSON:
     suggestedNextTest: "Try alternative payload or technique"
   };
 }
-function createSmartTestTool(session, model) {
+function createSmartTestTool(session, model, onTokenUsage) {
   return tool({
     name: "test_parameter",
     description: `Intelligently test a parameter for a vulnerability using AI-powered adaptive testing.
@@ -68830,7 +68914,7 @@ test_parameter({
           parameter,
           endpoint,
           context
-        }, model);
+        }, model, onTokenUsage);
         console.log(`Strategy: ${strategy}`);
         const results = [];
         let vulnerable = false;
@@ -68843,7 +68927,7 @@ test_parameter({
             context: { ...context, parameter, endpoint },
             previousResults: results,
             round
-          }, model);
+          }, model, onTokenUsage);
           console.log(`  Payload: ${payloadData.payload}`);
           console.log(`  Reasoning: ${payloadData.reasoning}`);
           let response;
@@ -68872,7 +68956,7 @@ test_parameter({
             attackType,
             knowledge,
             previousResults: results
-          }, model);
+          }, model, onTokenUsage);
           console.log(`  Analysis: ${analysis.reasoning}`);
           console.log(`  Vulnerable: ${analysis.vulnerable} (confidence: ${analysis.confidence})`);
           results.push({
@@ -69322,7 +69406,7 @@ function wrapCommandWithHeaders(command, headers) {
   }
   return wrapped;
 }
-function createPentestTools(session, model, toolOverride) {
+function createPentestTools(session, model, toolOverride, onTokenUsage) {
   const offensiveHeaders = getOffensiveHeaders(session);
   const rateLimiter = session._rateLimiter;
   const executeCommand = tool({
@@ -69496,7 +69580,7 @@ COMMON TESTING PATTERNS:
     http_request: httpRequest,
     document_finding: createDocumentFindingTool(session),
     record_test_result: createRecordTestResultTool(session),
-    test_parameter: createSmartTestTool(session, model || "claude-sonnet-4-20250514"),
+    test_parameter: createSmartTestTool(session, model || "claude-sonnet-4-20250514", onTokenUsage),
     check_testing_coverage: createCheckTestingCoverageTool(session),
     validate_completeness: createValidateCompletenessTool(session),
     enumerate_endpoints: createEnumerateEndpointsTool(session),
@@ -70555,6 +70639,7 @@ function runAgent(opts) {
     objective,
     model,
     onStepFinish,
+    onToolTokenUsage,
     abortSignal,
     silent,
     authConfig,
@@ -70574,7 +70659,7 @@ function runAgent(opts) {
     analyze_scan,
     scratchpad,
     generate_report
-  } = createPentestTools(session, undefined, toolOverride);
+  } = createPentestTools(session, undefined, toolOverride, onToolTokenUsage);
   const document_finding = tool({
     name: "document_finding",
     description: `Document a security finding with severity, impact, and remediation guidance.
@@ -72340,8 +72425,8 @@ function PentestAgentDisplay() {
   const { closePentest } = useCommand();
   const {
     model,
-    addTokens,
-    setTokenCount,
+    addTokenUsage,
+    resetTokenUsage,
     setThinking,
     isExecuting,
     setIsExecuting
@@ -72506,6 +72591,7 @@ function PentestAgentDisplay() {
       setHasStarted(true);
       setThinking(true);
       setIsExecuting(true);
+      resetTokenUsage();
       const controller = new AbortController;
       setAbortController(controller);
       try {
@@ -72516,8 +72602,12 @@ function PentestAgentDisplay() {
           abortSignal: controller.signal,
           sessionConfig,
           onStepFinish: ({ usage }) => {
-            const stepTokens = (usage.inputTokens ?? 0) + (usage.outputTokens ?? 0);
-            setTokenCount(stepTokens);
+            const inputTokens = usage.inputTokens ?? 0;
+            const outputTokens = usage.outputTokens ?? 0;
+            addTokenUsage(inputTokens, outputTokens);
+          },
+          onToolTokenUsage: (inputTokens, outputTokens) => {
+            addTokenUsage(inputTokens, outputTokens);
           }
         });
         setSessionPath(result.session.rootPath);
@@ -72535,7 +72625,6 @@ Path: ${result.session.rootPath}`,
         for await (const delta of result.fullStream) {
           if (delta.type === "text-delta") {
             currentAssistantText += delta.text;
-            addTokens(1);
             const lastMessage = allMessages[allMessages.length - 1];
             if (lastMessage && lastMessage.role === "assistant") {
               lastMessage.content = currentAssistantText;
@@ -74827,7 +74916,7 @@ If resuming from summarization, review the assets in the session assets folder a
 import { join as join5 } from "path";
 import { writeFileSync as writeFileSync5, mkdirSync as mkdirSync5, existsSync as existsSync10 } from "fs";
 function runAgent2(opts) {
-  const { target, model, onStepFinish, abortSignal } = opts;
+  const { target, model, onStepFinish, onToolTokenUsage, abortSignal } = opts;
   const session = opts.session || createSession(target);
   const subagentId = `attack-surface-${nanoid3(6)}`;
   console.log(`Created attack surface session: ${session.id}`);
@@ -74836,7 +74925,7 @@ function runAgent2(opts) {
   if (!existsSync10(assetsPath)) {
     mkdirSync5(assetsPath, { recursive: true });
   }
-  const { analyze_scan, execute_command, http_request } = createPentestTools(session, model);
+  const { analyze_scan, execute_command, http_request } = createPentestTools(session, model, undefined, onToolTokenUsage);
   const document_asset = tool({
     name: "document_asset",
     description: `Document a discovered asset during attack surface analysis.
@@ -75133,13 +75222,14 @@ function runAgent3(opts) {
     onSubagentSpawn,
     onSubagentMessage,
     onSubagentComplete,
+    onSubagentTokenUsage,
     session: sessionProp
   } = opts;
   const session = sessionProp || createSession(target, undefined, undefined, sessionConfig);
   const logger = new Logger(session);
   logger.log(`Created thorough pentest session: ${session.id}`);
   logger.log(`Session path: ${session.rootPath}`);
-  const tools2 = createOrchestratorTools(session, model, abortSignal, onSubagentSpawn, onSubagentMessage, onSubagentComplete, logger);
+  const tools2 = createOrchestratorTools(session, model, abortSignal, onSubagentSpawn, onSubagentMessage, onSubagentComplete, onSubagentTokenUsage, logger);
   const enhancedPrompt = `
 TARGET: ${target}
 
@@ -75175,7 +75265,7 @@ Begin by using the get_attack_surface tool to map the complete attack surface of
   streamResult.session = session;
   return { streamResult, session };
 }
-function createOrchestratorTools(session, model, abortSignal, onSubagentSpawn, onSubagentMessage, onSubagentComplete, logger) {
+function createOrchestratorTools(session, model, abortSignal, onSubagentSpawn, onSubagentMessage, onSubagentComplete, onSubagentTokenUsage, logger) {
   const getAttackSurface = tool({
     name: "get_attack_surface",
     description: `Run the attack surface analysis agent to discover all assets and identify targets.
@@ -75207,7 +75297,19 @@ Use this as the FIRST step in your thorough penetration test.`,
           target,
           objective,
           model,
-          abortSignal
+          abortSignal,
+          onStepFinish: ({ usage }) => {
+            if (onSubagentTokenUsage) {
+              const inputTokens = usage.inputTokens ?? 0;
+              const outputTokens = usage.outputTokens ?? 0;
+              onSubagentTokenUsage(subagentId, inputTokens, outputTokens);
+            }
+          },
+          onToolTokenUsage: (inputTokens, outputTokens) => {
+            if (onSubagentTokenUsage) {
+              onSubagentTokenUsage(subagentId, inputTokens, outputTokens);
+            }
+          }
         });
         const allMessages = [];
         let currentAssistantText = "";
@@ -75382,7 +75484,19 @@ You can spawn multiple agents in parallel - they will run concurrently.`,
               target: targetInfo.target,
               objective: targetInfo.objective,
               model,
-              abortSignal
+              abortSignal,
+              onStepFinish: ({ usage }) => {
+                if (onSubagentTokenUsage) {
+                  const inputTokens = usage.inputTokens ?? 0;
+                  const outputTokens = usage.outputTokens ?? 0;
+                  onSubagentTokenUsage(subagentId, inputTokens, outputTokens);
+                }
+              },
+              onToolTokenUsage: (inputTokens, outputTokens) => {
+                if (onSubagentTokenUsage) {
+                  onSubagentTokenUsage(subagentId, inputTokens, outputTokens);
+                }
+              }
             });
             const allMessages = [];
             let currentAssistantText = "";
@@ -75709,8 +75823,8 @@ function useThoroughPentestAgent() {
   const [abortController, setAbortController] = import_react22.useState(null);
   const {
     model,
-    addTokens,
-    setTokenCount,
+    addTokenUsage,
+    resetTokenUsage,
     setThinking,
     isExecuting,
     setIsExecuting
@@ -75738,6 +75852,7 @@ function useThoroughPentestAgent() {
       setHasStarted(true);
       setThinking(true);
       setIsExecuting(true);
+      resetTokenUsage();
       const controller = new AbortController;
       setAbortController(controller);
       try {
@@ -75747,8 +75862,12 @@ function useThoroughPentestAgent() {
           abortSignal: controller.signal,
           sessionConfig,
           onStepFinish: ({ usage }) => {
-            const stepTokens = (usage.inputTokens ?? 0) + (usage.outputTokens ?? 0);
-            setTokenCount(stepTokens);
+            const inputTokens = usage.inputTokens ?? 0;
+            const outputTokens = usage.outputTokens ?? 0;
+            addTokenUsage(inputTokens, outputTokens);
+          },
+          onSubagentTokenUsage: (subagentId, input, output) => {
+            addTokenUsage(input, output);
           },
           onSubagentSpawn: (subagentInfo) => {
             setSubagents((prev) => [
@@ -75859,7 +75978,6 @@ Mode: Pentest (Orchestrator)`,
         await consumeStream2(result, {
           onTextDelta: (delta) => {
             currentAssistantText += delta.text;
-            addTokens(1);
             const lastMessage = allMessages[allMessages.length - 1];
             if (lastMessage && lastMessage.role === "assistant") {
               lastMessage.content = currentAssistantText;