@pensar/apex 2.6.1 → 2.7.0-canary.1572d9a4

package/src/core/agents/offSecAgent/offensiveSecurityAgent.ts

@@ -8,7 +8,7 @@ import type {
   TextStreamPart,
   ToolSet,
 } from "ai";
-import { hasToolCall } from "ai";
+import { hasToolCall, stepCountIs } from "ai";
 import { normalizeStepUsage, streamResponse } from "../../ai";
 import { AgentEventBus, type StreamIdContext } from "../../eventBus";
 import { newMessageId, newPartId } from "../../id/id";
@@ -25,6 +25,8 @@ import { ApprovalDeniedError } from "../../operator";
 import type { RunContextRecorder } from "../../runtime/runContext";
 import { create as createSession, type SessionInfo } from "../../session";
 import { scopedLogger } from "../../util/lazyLogger";
+import { LANCE_CONFIG, LANCE_SYSTEM_PROMPT, LanceRuntime } from "../lance";
+import { createLanceJudgeVerifier } from "../lanceJudge";
 import {
   createInterruptedStepFinalizer,
   type FinalizeInterruptedStepInput,
@@ -43,6 +45,7 @@ import {
 } from "./runtimeContext";
 import { responseArgBytes, StreamDiagnostics } from "./streamDiagnostics";
 import { inProcessSubagentSpawner } from "./subagentSpawner";
+import { wrapToolExecutionContext } from "./toolExecutionContext";
 import { ToolLifecycleTracker } from "./toolLifecycle";
 import {
   ASK_USER_QUESTIONS_TOOL_NAME,
@@ -298,6 +301,7 @@ export class OffensiveSecurityAgent<TResult = void> {
 
   private readonly resolveResult?: (
     streamResult: StreamTextResult<ToolSet, never>,
+    validatedResponse?: unknown,
   ) => TResult | Promise<TResult>;
 
   /** The structured result captured by the `response` tool, or `null` until it fires. */
@@ -423,8 +427,18 @@ export class OffensiveSecurityAgent<TResult = void> {
     let session = input.session;
     if (!session) {
       session = await createSession({
+        ...(input.mode === "lance" ? { name: "Lance" } : {}),
         targets: input.target ? [input.target] : [],
-        config: input.sessionConfig,
+        config:
+          input.reasoning === undefined
+            ? input.sessionConfig
+            : {
+                ...input.sessionConfig,
+                modelReasoning: {
+                  model: input.model,
+                  setting: input.reasoning,
+                },
+              },
         model: input.model,
         authConfig: input.authConfig,
         userMessage: input.prompt,
@@ -444,7 +458,65 @@ export class OffensiveSecurityAgent<TResult = void> {
     return this.subagentId ?? this._session.id;
   }
 
+  readonly lanceRuntime?: LanceRuntime;
+  private lanceModel?: string;
+
   constructor(input: OffensiveSecurityAgentInput<TResult>) {
+    if (
+      input.mode === "lance" &&
+      !input.lance?.interactive &&
+      existsSync(join(input.session.rootPath, "lance-result.json"))
+    ) {
+      throw new Error("This Lance attempt is terminal; use a new session");
+    }
+    if (input.mode === "lance") {
+      if (
+        input.extraTools ||
+        input.responseSchema ||
+        input.subagentSpawner ||
+        input.system ||
+        input.resolveResult ||
+        input.subagentId
+      ) {
+        throw new Error(
+          "Lance does not permit tool, prompt, response, or spawner overrides",
+        );
+      }
+      const lanceOptions = { ...input.lance };
+      this.lanceRuntime = new LanceRuntime(
+        lanceOptions,
+        input.abortSignal,
+        input.session.rootPath,
+      );
+      if (input.findingPolicy === "native" && !lanceOptions.verifyObjective) {
+        lanceOptions.verificationSource = "apex_judge";
+        lanceOptions.verifyObjective = createLanceJudgeVerifier(
+          input,
+          this.lanceRuntime,
+        );
+      }
+      this.lanceModel = input.model;
+      const middleware = input.languageModelMiddleware;
+      input = {
+        ...input,
+        system: LANCE_SYSTEM_PROMPT,
+        abortSignal: this.lanceRuntime.signal,
+        subagentSpawner: this.lanceRuntime.spawner,
+        languageModelMiddleware: [
+          this.lanceRuntime.middleware,
+          ...(Array.isArray(middleware)
+            ? middleware
+            : middleware
+              ? [middleware]
+              : []),
+        ],
+        stopWhen: [
+          () => this.lanceRuntime!.terminal,
+          stepCountIs(LANCE_CONFIG.maxSteps),
+        ],
+        skillsRegistry: undefined,
+      };
+    }
     this._session = input.session;
     this.subagentId = input.subagentId;
     this.subagentName = input.subagentName;
@@ -552,6 +624,7 @@ export class OffensiveSecurityAgent<TResult = void> {
       input.credentialManager ?? input.session.credentialManager;
 
     const toolCtx = {
+      findingPolicy: input.findingPolicy,
       ...(input.attackSurfaceArtifactsPath
         ? { attackSurfaceArtifactsPath: input.attackSurfaceArtifactsPath }
         : {}),
@@ -589,6 +662,7 @@ export class OffensiveSecurityAgent<TResult = void> {
       enableThinking: input.enableThinking,
       thinkingEffort: input.thinkingEffort,
       openAIReasoningEffort: input.openAIReasoningEffort,
+      reasoning: input.reasoning,
       surfaceIntegrationEnabled: input.surfaceIntegrationEnabled,
       projectThreatModel: input.projectThreatModel,
       planSubagentId: input.planSubagentId,
@@ -607,6 +681,7 @@ export class OffensiveSecurityAgent<TResult = void> {
       emailAdapterFor: input.emailAdapterFor,
       languageModelMiddleware: input.languageModelMiddleware,
       usageRecorder: input.usageRecorder,
+      toolExecutionContext: input.toolExecutionContext,
       streamIdFactory: input.streamIdFactory,
     };
 
@@ -630,7 +705,9 @@ export class OffensiveSecurityAgent<TResult = void> {
     let activeTools = (input.activeTools as string[]).filter(
       passesEmailSmsGates,
     );
-    if (input.mode === "plan") {
+    if (input.mode === "lance") {
+      activeTools = [...LANCE_CONFIG.tools];
+    } else if (input.mode === "plan") {
       const planSet = new Set<string>(PLAN_MODE_TOOL_NAMES);
       activeTools = activeTools.filter((t) => planSet.has(t));
     } else if (input.mode === "fast-strike") {
@@ -658,6 +735,12 @@ export class OffensiveSecurityAgent<TResult = void> {
       }
     }
 
+    if (input.findingPolicy === "external") {
+      activeTools = activeTools.filter(
+        (name) => name !== "document_vulnerability",
+      );
+    }
+
     activeTools = filterWorkspaceToolsForRun(
       activeTools,
       input.prompt,
@@ -666,10 +749,37 @@ export class OffensiveSecurityAgent<TResult = void> {
 
     const builtinTools = createToolsForNames(toolCtx, activeTools);
 
-    let tools: ToolSet = input.extraTools
-      ? { ...builtinTools, ...input.extraTools }
+    const extraTools =
+      input.findingPolicy === "external" && input.extraTools
+        ? Object.fromEntries(
+            Object.entries(input.extraTools).filter(
+              ([name]) => name !== "document_vulnerability",
+            ),
+          )
+        : input.extraTools;
+    let tools: ToolSet = extraTools
+      ? { ...builtinTools, ...extraTools }
       : { ...builtinTools };
 
+    if (this.lanceRuntime) {
+      const runtime = this.lanceRuntime;
+      tools = Object.fromEntries(
+        Object.entries(tools).map(([name, definition]) => [
+          name,
+          {
+            ...definition,
+            execute: async (
+              args: unknown,
+              options: Parameters<NonNullable<typeof definition.execute>>[1],
+            ) => {
+              runtime.checkAgent();
+              return definition.execute!(args, options);
+            },
+          },
+        ]),
+      );
+    }
+
     // Only approved calls may reach the execution journal.
     if (input.toolExecutionRecorder) {
       tools = wrapRecordedTools(tools, input.toolExecutionRecorder);
@@ -681,9 +791,15 @@ export class OffensiveSecurityAgent<TResult = void> {
         tools,
         input.approvalGate,
         AGENT_PAUSE_TOOLS,
+        this.lanceRuntime?.signal,
       );
     }
 
+    if (this.lanceRuntime) {
+      tools = this.lanceRuntime.tools(tools);
+      activeTools = Object.keys(tools);
+    }
+
     // -- Response schema → auto capture / stop / resolve ----------------------
     if (input.responseSchema) {
       tools = {
@@ -712,7 +828,7 @@ export class OffensiveSecurityAgent<TResult = void> {
                         Promise.race([sr.response, this._drainRejection]),
                     },
                   });
-                  return customResolve(guarded);
+                  return customResolve(guarded, result);
                 })
                 .then((resolved) => {
                   this._capturedResponse = resolved;
@@ -730,6 +846,10 @@ export class OffensiveSecurityAgent<TResult = void> {
       };
     }
 
+    if (input.toolExecutionContext) {
+      tools = wrapToolExecutionContext(tools, input.toolExecutionContext);
+    }
+
     if (input.resolveResult) {
       this.resolveResult = input.resolveResult;
     } else if (input.responseSchema) {
@@ -759,16 +879,36 @@ export class OffensiveSecurityAgent<TResult = void> {
     }
     this.messagesPath = join(messagesDir, "messages.json");
 
-    // Mutable so that summarization can clear stale history.
-    const initialMessagesRef: { current: ModelMessage[] } = {
-      current: input.messages
-        ? [...input.messages]
-        : [
-            {
-              role: "user" as const,
-              content: [{ type: "text", text: input.prompt }],
-            },
-          ],
+    const preparedContextRef = {
+      current: {
+        messages: input.messages
+          ? [...input.messages]
+          : [
+              {
+                role: "user" as const,
+                content: [{ type: "text" as const, text: input.prompt }],
+              },
+            ],
+        responseMessageOffset: 0,
+      },
+    };
+    const legacyMessages = (responseMessages: ModelMessage[]) => [
+      ...preparedContextRef.current.messages,
+      ...responseMessages.slice(
+        preparedContextRef.current.responseMessageOffset,
+      ),
+    ];
+    const persistedMessages = (messages: ModelMessage[]) => {
+      // Cache-marked system messages are transport context; the next turn
+      // supplies a fresh system prompt and must not replay saved copies.
+      let first = 0;
+      while (
+        messages[first]?.role === "system" &&
+        messages[first]?.providerOptions
+      ) {
+        first++;
+      }
+      return first === 0 ? messages : messages.slice(first);
     };
 
     // Agent-local persistence: debounce, latest snapshot, and the
@@ -796,6 +936,7 @@ export class OffensiveSecurityAgent<TResult = void> {
       input.system ??
       buildBaseSystemPrompt({
         sandboxMode: agentCwd === input.session.rootPath,
+        findingPolicy: input.findingPolicy,
       });
     const commandPlatform = resolveCommandPlatform(toolCtx);
     const workspaceSection = buildSessionWorkspaceSection(
@@ -803,6 +944,7 @@ export class OffensiveSecurityAgent<TResult = void> {
       agentCwd,
       activeTools,
       fileWorkspaceRoot,
+      input.findingPolicy,
     );
     // Bundled wordlist paths are host-local: advertise them only when
     // commands execute locally.
@@ -824,9 +966,9 @@ export class OffensiveSecurityAgent<TResult = void> {
     // Only agents whose tools can execute commands, touch files, or run
     // PoCs get a runtime probe — for anyone else the facts are noise and the
     // probe is a wasted subprocess.
-    this.probesRuntime = activeTools.some((name) =>
-      RUNTIME_CAPABLE_TOOL_NAMES.has(name),
-    );
+    this.probesRuntime =
+      input.mode !== "lance" &&
+      activeTools.some((name) => RUNTIME_CAPABLE_TOOL_NAMES.has(name));
 
     // -- Init record + stream --------------------------------------------------
     // Created by async initialization (streamReady) so the AI SDK telemetry
@@ -840,7 +982,10 @@ export class OffensiveSecurityAgent<TResult = void> {
       // so the hash is stable across runs with identical prompt versions. No
       // step record can precede this init line: steps only originate from the
       // stream created below.
-      const baseWithFacts = assembleBaseSystemPrompt(facts);
+      const baseWithFacts =
+        input.mode === "lance"
+          ? LANCE_SYSTEM_PROMPT
+          : assembleBaseSystemPrompt(facts);
       traceWriter.writeInit({
         model: input.model,
         systemPrompt: baseWithFacts,
@@ -848,11 +993,30 @@ export class OffensiveSecurityAgent<TResult = void> {
         sessionId: input.session.id,
         target: input.target,
       });
+      this.lanceRuntime?.start();
+      const continuation = this.lanceRuntime?.continuationContext();
       return streamResponse({
-        prompt: input.prompt,
-        system: baseWithFacts + workspaceSection,
+        prompt: continuation
+          ? `${input.prompt}\n\n${continuation}`
+          : input.prompt,
+        system:
+          input.mode === "lance"
+            ? `${LANCE_SYSTEM_PROMPT}\nWorking directory (not automatically in scope): ${agentCwd}\nTarget: ${input.target ?? "Use the user's explicit target, or ask if missing"}\nDeadline: ${new Date(this.lanceRuntime!.softDeadlineAt).toISOString()}`
+            : baseWithFacts + workspaceSection,
+        primaryModelOnly: input.mode === "lance",
+        _initialPrompt: input.mode === "lance" ? input.prompt : undefined,
+        runContextCompaction: this.lanceRuntime?.runContextCompaction,
+        getContinuationContext: this.lanceRuntime?.continuationContext,
+        onInvalidToolCall: (name) => {
+          if (this.lanceRuntime && !activeTools.includes(name)) {
+            this.lanceRuntime.violation("forbidden_tool", name);
+          }
+        },
         model: input.model,
-        messages: input.messages,
+        messages:
+          continuation && input.messages
+            ? [...input.messages, { role: "user", content: continuation }]
+            : input.messages,
         tools,
         activeTools,
         stopWhen,
@@ -861,6 +1025,15 @@ export class OffensiveSecurityAgent<TResult = void> {
         usageRecorder: input.usageRecorder,
         contextRecorder: input.contextRecorder,
         inferenceRecorder: input.inferenceRecorder,
+        onContextPrepared: (context) => {
+          preparedContextRef.current = {
+            messages: [...context.messages],
+            responseMessageOffset: context.responseMessageOffset,
+          };
+          if (this.lanceRuntime) {
+            this.writer.setLatest(persistedMessages(context.messages));
+          }
+        },
         // Per-subagent so the overflow tool-result dumps land next to this
         // agent's messages.json (`subagents/{id}/tool-results/`) and a host
         // can reclaim them when the subagent finishes, instead of piling up
@@ -880,16 +1053,11 @@ export class OffensiveSecurityAgent<TResult = void> {
             return;
           }
 
-          // Recorded runs project the canonical committed context; the
-          // legacy path reconstructs from the (possibly stale after a fit)
-          // initial prefix. Canonical latest() already includes this step's
-          // messages — committing happens upstream in the ai layer's
-          // onStepFinish wrapper, before this user callback runs.
-          const canonical = input.contextRecorder?.latest() ?? [
-            ...initialMessagesRef.current,
-            ...event.response.messages,
-          ];
-          this.writer.setLatest(canonical);
+          this.lanceRuntime?.stepFinished();
+          const canonical =
+            input.contextRecorder?.latest() ??
+            legacyMessages(event.response.messages);
+          this.writer.setLatest(persistedMessages(canonical));
           schedulePersist();
           traceWriter.recordStep(
             event.response.messages as ModelMessage[],
@@ -905,9 +1073,10 @@ export class OffensiveSecurityAgent<TResult = void> {
           await input.onStepFinish?.(event);
         },
         onSummarized: () => {
-          // Context was reset by summarization — discard the old history so
-          // subsequent onStepFinish writes only persist post-summary messages.
-          initialMessagesRef.current = [];
+          preparedContextRef.current = {
+            messages: [],
+            responseMessageOffset: 0,
+          };
           traceWriter.markSummarized();
         },
         onFinish: async (event) => {
@@ -915,11 +1084,11 @@ export class OffensiveSecurityAgent<TResult = void> {
           this.writer.cancelTimer();
           // Skip if the interrupted-step finalizer already wrote the abort snapshot.
           if (!this.writer.syntheticsPersisted) {
-            const finalMessages = input.contextRecorder?.latest() ??
-              this.writer.latest ?? [
-                ...initialMessagesRef.current,
-                ...event.response.messages,
-              ];
+            const finalMessages = persistedMessages(
+              input.contextRecorder?.latest() ??
+                this.writer.latest ??
+                legacyMessages(event.response.messages),
+            );
             this.writer.setLatest(finalMessages);
             await this.writer.enqueueWrite(finalMessages).catch(() => {});
           }
@@ -931,6 +1100,7 @@ export class OffensiveSecurityAgent<TResult = void> {
         enableThinking: input.enableThinking,
         thinkingEffort: input.thinkingEffort,
         openAIReasoningEffort: input.openAIReasoningEffort,
+        reasoning: input.reasoning,
         silent: true,
       });
     };
@@ -1039,6 +1209,21 @@ export class OffensiveSecurityAgent<TResult = void> {
    * @returns The value produced by `resolveResult`, or `void` if none was provided.
    */
   async consume(): Promise<TResult> {
+    if (!this.lanceRuntime) return this.consumeRun();
+    this.lanceRuntime.start();
+    let error: unknown;
+    try {
+      return await this.consumeRun();
+    } catch (cause) {
+      error = cause;
+      return undefined as TResult;
+    } finally {
+      this.lanceRuntime.finish(error);
+      this.lanceRuntime.persist(this.session.rootPath, this.lanceModel!);
+    }
+  }
+
+  private async consumeRun(): Promise<TResult> {
     const sid = this.subagentId;
     const bus = this.eventBus;
 
@@ -1196,6 +1381,13 @@ export class OffensiveSecurityAgent<TResult = void> {
             }
             return result;
           } catch (err) {
+            if (this.lanceRuntime?.terminal) {
+              span.setAttribute(
+                "pensar.lance.outcome",
+                this.lanceRuntime.metrics(this.lanceModel!).outcome,
+              );
+              throw err;
+            }
             // Record once, keep the cardinality low, preserve the original.
             span.recordException(err as Error);
             span.setStatus({ code: SpanStatusCode.ERROR });
@@ -1550,6 +1742,7 @@ function wrapToolsWithApprovalGate(
   tools: ToolSet,
   gate: ApprovalGate,
   exemptToolNames?: Set<string>,
+  signal?: AbortSignal,
 ): ToolSet {
   const wrapped: ToolSet = {};
 
@@ -1578,7 +1771,15 @@ function wrapToolsWithApprovalGate(
           `tc_${Date.now()}_${Math.random().toString(36).slice(2, 8)}`;
 
         log.debug(`${name} (${toolCallId}): checking`);
+        const cancelApproval = () => {
+          for (const pending of gate.getPendingApprovals()) {
+            if (pending.toolCallId === String(toolCallId))
+              gate.deny(pending.id);
+          }
+        };
         try {
+          signal?.throwIfAborted();
+          signal?.addEventListener("abort", cancelApproval, { once: true });
           await gate.check(name, String(toolCallId), args);
         } catch (err) {
           if (err instanceof ApprovalDeniedError) {
@@ -1586,6 +1787,8 @@ function wrapToolsWithApprovalGate(
             return { blocked: true, reason: "Denied by operator" };
           }
           throw err;
+        } finally {
+          signal?.removeEventListener("abort", cancelApproval);
         }
         log.debug(`${name} (${toolCallId}): approved, executing`);
 

package/src/core/agents/offSecAgent/prompt.test.ts

@@ -70,3 +70,49 @@ describe("buildBaseSystemPrompt workspace management", () => {
     );
   });
 });
+
+describe("external finding policy prompts", () => {
+  const session = {
+    rootPath: "/tmp/apex-external-prompt",
+    findingsPath: "/tmp/apex-external-prompt/findings",
+    pocsPath: "/tmp/apex-external-prompt/pocs",
+    scratchpadPath: "/tmp/apex-external-prompt/scratchpad",
+    logsPath: "/tmp/apex-external-prompt/logs",
+  };
+
+  it("keeps exploration guidance without directing an unavailable finding tool", () => {
+    const prompt = buildBaseSystemPrompt({ findingPolicy: "external" });
+    expect(prompt).toContain("Source code analysis");
+    expect(prompt).toContain("caller grades them after the run");
+    expect(prompt).not.toContain("document_vulnerability");
+    expect(prompt).not.toContain("published PoCs");
+  });
+
+  it.each([
+    [session.rootPath, undefined],
+    ["/project", undefined],
+    ["/project", "/tmp/apex-external-prompt/helpers"],
+  ])("uses the caller's finding contract for cwd %s", (cwd, fileWorkspaceRoot) => {
+    const prompt = buildSessionWorkspaceSection(
+      session,
+      cwd,
+      ["execute_command", "read_file"],
+      fileWorkspaceRoot,
+      "external",
+    );
+    expect(prompt).toContain("caller's output contract");
+    expect(prompt).not.toContain("document_vulnerability");
+  });
+
+  it("keeps source guidance policy-safe even when tool names are not supplied", () => {
+    const prompt = buildSessionWorkspaceSection(
+      { ...session, config: { codebasePath: "/project" } },
+      "/project",
+      undefined,
+      undefined,
+      "external",
+    );
+    expect(prompt).toContain("Source Code Assessment");
+    expect(prompt).not.toContain("document_vulnerability");
+  });
+});