@gpzhang2001/sharpkit-team 0.2.1

package/src/index.ts

@@ -0,0 +1,427 @@
+/**
+ * Red-team orchestration — the M4 tool-team layer: strix-named tools
+ * (create_agent / send_message_to_agent / wait_for_agents /
+ * view_agent_graph / stop_agent / agent_finish) composed over the dsh
+ * subagent SERVICE (ctx.subagents + the spawn provider's continuable
+ * children), not over tools. Children inherit the parent Agent's preset and
+ * full tool surface (in-process spawn semantics); requested skills are
+ * injected as a directive line in the child's prompt. S2-verified semantics
+ * hold: send steers at step boundaries, interrupt parks the inbox (children
+ * stay resumable), completion notices reach the parent's turn boundary.
+ * A token-budget circuit breaker interrupts every tracked child and blocks
+ * new spawns when the session's accumulated usage crosses the ceiling.
+ * @module @gpzhang2001/sharpkit-team
+ */
+
+import type { Context } from '@deepseek-ai/cordis'
+import type Schema from '@deepseek-ai/schemastery'
+import z from '@deepseek-ai/schemastery'
+import { defineTool } from '@deepseek-ai/dsh-tools'
+import type { ContentBlock } from '@deepseek-ai/dsh-llm'
+import type {} from '@deepseek-ai/dsh-subagent'
+import type {} from '@deepseek-ai/dsh-session'
+
+/** Structural view of the subagent service this package drives. */
+export interface SubagentsLike {
+  start(name: string, request: {
+    readonly label?: string
+    readonly prompt: ContentBlock[]
+    readonly parent: unknown
+    readonly signal: AbortSignal
+  }): { readonly id: { readonly [key: string]: unknown } | string; readonly result: Promise<{ readonly stopReason?: string; readonly output?: readonly ContentBlock[] }> }
+  sendMessage(sender: unknown, targetId: unknown, content: ContentBlock[], options?: unknown): Promise<unknown>
+  interrupt(targetSessionId: unknown, authority:
+    | { readonly kind: 'user'; readonly parentSessionId: unknown }
+    | { readonly kind: 'ancestor'; readonly agent: unknown }): void
+}
+
+/** Deployment-tunable configuration. */
+export interface Config {
+  /** Maximum delegation depth (root=1 spawns children; children don't spawn). */
+  readonly maxTeamDepth?: number
+  /** Explicit session token ceiling for the circuit breaker (overrides budget estimates). */
+  readonly maxSessionTokens?: number
+  /** USD budget ceiling for the circuit breaker (token-estimated; see usdPerMillionTokens). */
+  readonly maxBudgetUsd?: number
+  /** Estimated blended $/1M tokens mapping a USD budget to a token ceiling. */
+  readonly usdPerMillionTokens?: number
+  /** Skills catalog root — only used to validate requested skill names exist. */
+  readonly skillsRoot?: string
+}
+
+export const name = 'pentest-tool-team'
+
+export const inject = ['tools']
+
+export const Config: Schema<Config> = z.object({
+  maxTeamDepth: z.number().default(2),
+  maxSessionTokens: z.number(),
+  maxBudgetUsd: z.number(),
+  usdPerMillionTokens: z.number(),
+  skillsRoot: z.string(),
+})
+
+/** Default blended $/1M tokens when neither config nor preset supplies a mapping. */
+const DEFAULT_USD_PER_MILLION_TOKENS = 0.5
+
+/** Fallback token ceiling when no budget source is configured. */
+const DEFAULT_MAX_SESSION_TOKENS = 2_000_000
+
+/** One tracked child (roster row). */
+interface TrackedChild {
+  readonly id: string
+  readonly label: string
+  readonly task: string
+  readonly skills: readonly string[]
+  readonly startedAt: string
+  status: 'running' | 'completed' | 'stopped' | 'failed'
+  completionReport: string | undefined
+  readonly result: Promise<{ readonly stopReason?: string; readonly output?: readonly ContentBlock[] }>
+}
+
+/** The orchestrator-facing team handle (tests + UI consume). */
+export interface TeamHandle {
+  children(): ReadonlyArray<Readonly<TrackedChild>>
+  /** True once the circuit breaker tripped. */
+  isBreached(): boolean
+  tokensUsed(): number
+}
+
+/** Brand-free id extraction from the runtime's branded session ids. */
+function idText(value: unknown): string {
+  return typeof value === 'string' ? value : String(value)
+}
+
+export function apply(ctx: Context, config: Config = {}): TeamHandle {
+  const children = new Map<string, TrackedChild>()
+  /** Sessions whose usage counts toward this team's budget (root + children). */
+  const teamSessions = new Set<string>()
+  /** The first caller's Agent — the ancestor authority for breaker interrupts. */
+  let rootAgent: unknown
+  let tokensUsed = 0
+  let breached = false
+  const maxTeamDepth = config.maxTeamDepth ?? 2
+  const usdPerMillionTokens = config.usdPerMillionTokens ?? DEFAULT_USD_PER_MILLION_TOKENS
+
+  const subagents = (): SubagentsLike => ctx.subagents as unknown as SubagentsLike
+
+  /** Budget ceiling in tokens: explicit tokens > config USD > preset USD > fallback. */
+  const tokenCeiling = (): number => {
+    if (config.maxSessionTokens !== undefined) return config.maxSessionTokens
+    const preset = ctx.get('pentestPreset') as { maxBudgetUsd?: number } | undefined
+    const budgetUsd = config.maxBudgetUsd ?? preset?.maxBudgetUsd
+    if (budgetUsd !== undefined) return Math.max(1, Math.floor((budgetUsd / usdPerMillionTokens) * 1_000_000))
+    return DEFAULT_MAX_SESSION_TOKENS
+  }
+
+  // Budget circuit breaker: accumulate this team's session usage and trip the
+  // breaker. The interrupt goes out under ancestor authority from the root
+  // agent (kind:'user' requires a human-presented parent session id, which a
+  // plugin-side breaker does not have).
+  void ctx.on('session/event', (session: { id?: unknown }, event: unknown) => {
+    const record = event as { type?: string; data?: { usage?: { inputTokens?: number; outputTokens?: number; input?: number; output?: number } } }
+    if (record.type !== 'assistant/message') return
+    const sessionId = session.id === undefined ? '' : String(session.id)
+    if (!teamSessions.has(sessionId)) return
+    const usage = record.data?.usage
+    if (usage === undefined || usage === null) return
+    const input = usage.inputTokens ?? usage.input ?? 0
+    const output = usage.outputTokens ?? usage.output ?? 0
+    if (typeof input === 'number' && typeof output === 'number') tokensUsed += input + output
+    if (!breached && tokensUsed > tokenCeiling()) {
+      breached = true
+      ctx.logger.warn(`pentest-team: session token budget exceeded (${String(tokensUsed)} > ${String(tokenCeiling())} tokens); interrupting ${String(children.size)} child agent(s)`)
+      for (const child of children.values()) {
+        if (child.status !== 'running') continue
+        if (rootAgent !== undefined) {
+          try {
+            subagents().interrupt(child.id as never, { kind: 'ancestor', agent: rootAgent })
+          } catch {
+            // Best-effort interruption; the child result settles the status.
+          }
+        }
+        child.status = 'stopped'
+      }
+    }
+  })
+
+  const handle: TeamHandle = {
+    children: () => [...children.values()],
+    isBreached: () => breached,
+    tokensUsed: () => tokensUsed,
+  }
+  ctx.provide('pentestTeam', handle)
+
+  const textBlock = (text: string): ContentBlock => ({ type: 'text', text })
+
+  ctx.tools.register(defineTool({
+    name: 'create_agent',
+    description: `Delegate a focused subtask to a specialist child agent. The child runs with this scan's preset and full tool surface. Pass skills (max 5, category/name) to point it at specialist knowledge. Do not run hands-on tests yourself that a child should run.`,
+    parameters: {
+      name: { type: 'string', required: true, description: 'Short specialist label (e.g. "recon", "web-sqli").' },
+      task: { type: 'string', required: true, description: 'The complete, self-contained task for the child.' },
+      skills: { type: 'array', items: { type: 'string' }, description: 'Up to 5 specialist skills (category/name) injected into the brief.' },
+    },
+    output: {
+      schema: {
+        type: 'object',
+        properties: {
+          success: { type: 'boolean', required: true },
+          agent_id: { type: 'string' },
+          status: { type: 'string' },
+          error: { type: 'string' },
+        },
+        additionalProperties: false,
+      },
+      render: (_args, value) => {
+        const result = value as { success: boolean; agent_id?: string; error?: string }
+        if (!result.success) return [{ type: 'text', text: `create_agent failed: ${result.error ?? 'unknown'}` }]
+        return [{ type: 'text', text: `child ${String(result.agent_id)} started` }]
+      },
+    },
+    execute: async (args, exec) => {
+      if (breached) return { success: false, error: `session token budget exhausted (${String(tokensUsed)} tokens) — wrap up and finish instead of spawning new agents` }
+      const runningCount = [...children.values()].filter(child => child.status === 'running').length
+      if (runningCount >= maxTeamDepth + 3) {
+        return { success: false, error: `too many concurrent children (${String(runningCount)}); wait_for_agents or stop_agent first` }
+      }
+      rootAgent ??= exec.agent
+      if (exec.agent !== undefined) teamSessions.add(String(exec.agent.session.id))
+      const brief: string[] = [`You are specialist agent "${args.name}" in an authorized penetration test.`, ``, `TASK:`, args.task]
+      if (args.skills !== undefined && args.skills.length > 0) {
+        brief.push('', 'SPECIALIST KNOWLEDGE: consult these skill areas and follow them.', ...args.skills.map(skill => `- ${skill}`))
+      }
+      brief.push('', 'Work autonomously. Report findings via create_vulnerability_report / create_dependency_report; record coverage with record_coverage; finish with a concise completion report as your final message.')
+      try {
+        const run = subagents().start('spawn', {
+          label: args.name,
+          prompt: [textBlock(brief.join('\n'))],
+          parent: exec.agent,
+          signal: exec.signal,
+        })
+        const id = idText(run.id)
+        // run.id is the child session id: its usage joins the team budget.
+        teamSessions.add(id)
+        const child: TrackedChild = {
+          id,
+          label: args.name,
+          task: args.task,
+          skills: args.skills ?? [],
+          startedAt: new Date().toISOString(),
+          status: 'running',
+          completionReport: undefined,
+          result: run.result,
+        }
+        children.set(id, child)
+        void run.result.then(outcome => {
+          if (child.status === 'running') {
+            child.status = outcome?.stopReason === 'error' ? 'failed' : 'completed'
+          }
+          // Fallback completion report: the child's final assistant text.
+          if (child.completionReport === undefined && outcome?.output !== undefined) {
+            const text = outcome.output.map(block => block.type === 'text' ? block.text : '').filter(part => part !== '').join('\n')
+            if (text !== '') child.completionReport = text
+          }
+        }, () => {
+          if (child.status === 'running') child.status = 'failed'
+        })
+        return { success: true, agent_id: id, status: 'running' }
+      } catch (error) {
+        return { success: false, error: String(error instanceof Error ? error.message : error) }
+      }
+    },
+  }))
+
+  ctx.tools.register(defineTool({
+    name: 'send_message_to_agent',
+    description: 'Steer a running child mid-run: new information, a course correction, or a request to wrap up. The message is delivered at the child\'s next step boundary.',
+    parameters: {
+      agent_id: { type: 'string', required: true, description: 'Child id from create_agent.' },
+      message: { type: 'string', required: true, description: 'The message text.' },
+    },
+    output: {
+      schema: {
+        type: 'object',
+        properties: {
+          success: { type: 'boolean', required: true },
+          error: { type: 'string' },
+        },
+        additionalProperties: false,
+      },
+      render: (_args, value) => {
+        const result = value as { success: boolean; error?: string }
+        return [{ type: 'text', text: result.success ? 'message delivered' : `send_message_to_agent failed: ${result.error ?? 'unknown'}` }]
+      },
+    },
+    execute: async (args, exec) => {
+      const child = children.get(args.agent_id)
+      if (child === undefined) return { success: false, error: `No child agent '${args.agent_id}'. Known: ${[...children.keys()].join(', ') || 'none'}.` }
+      try {
+        await subagents().sendMessage(exec.agent, args.agent_id as never, [textBlock(args.message)])
+        return { success: true }
+      } catch (error) {
+        return { success: false, error: String(error instanceof Error ? error.message : error) }
+      }
+    },
+  }))
+
+  ctx.tools.register(defineTool({
+    name: 'wait_for_agents',
+    description: 'Block until the named children report back (settlement notices also arrive automatically). Issue exactly ONE wait and react to what it returns.',
+    parameters: {
+      agent_ids: { type: 'array', items: { type: 'string' }, required: true, description: 'Child ids to wait for.' },
+    },
+    output: {
+      schema: {
+        type: 'object',
+        properties: {
+          success: { type: 'boolean', required: true },
+          agents: {
+            type: 'array',
+            required: true,
+            items: { type: 'object', properties: { agent_id: { type: 'string' }, status: { type: 'string' }, completion_report: { type: 'string' } }, additionalProperties: false },
+          },
+          error: { type: 'string' },
+        },
+        additionalProperties: false,
+      },
+      render: (_args, value) => {
+        const result = value as { agents?: Array<{ status: string }> }
+        return [{ type: 'text', text: `${String(result.agents?.length ?? 0)} child agent(s) settled` }]
+      },
+    },
+    execute: async args => {
+      const tracked = args.agent_ids.map(id => children.get(id)).filter((child): child is TrackedChild => child !== undefined)
+      if (tracked.length === 0) {
+        return { success: false, agents: [], error: `No known children among: ${args.agent_ids.join(', ')}` }
+      }
+      await Promise.all(tracked.map(child => child.result.catch(() => undefined)))
+      return {
+        success: true,
+        agents: tracked.map(child => ({
+          agent_id: child.id,
+          status: child.status,
+          ...(child.completionReport !== undefined ? { completion_report: child.completionReport } : {}),
+        })),
+      }
+    },
+  }))
+
+  ctx.tools.register(defineTool({
+    name: 'view_agent_graph',
+    description: 'Your live map of the team: every child agent with its id, label, task, skills, and status. Call it before spawning (avoid duplicates) and before finishing (no child still running).',
+    parameters: {},
+    output: {
+      schema: {
+        type: 'object',
+        properties: {
+          success: { type: 'boolean', required: true },
+          agents: {
+            type: 'array',
+            required: true,
+            items: {
+              type: 'object',
+              properties: {
+                agent_id: { type: 'string', required: true },
+                name: { type: 'string' },
+                task: { type: 'string' },
+                skills: { type: 'array', items: { type: 'string' } },
+                status: { type: 'string', required: true },
+                started_at: { type: 'string' },
+              },
+              additionalProperties: false,
+            },
+          },
+          tokens_used: { type: 'integer' },
+          budget_breached: { type: 'boolean' },
+        },
+        additionalProperties: false,
+      },
+      render: (_args, value) => {
+        const result = value as { agents: Array<{ name?: string; status: string }>; tokens_used: number }
+        return [{ type: 'text', text: `${String(result.agents.length)} child agent(s), ${String(result.tokens_used)} tokens used` }]
+      },
+    },
+    execute: async () => ({
+      success: true,
+      agents: [...children.values()].map(child => ({
+        agent_id: child.id,
+        name: child.label,
+        task: child.task,
+        skills: [...child.skills],
+        status: child.status,
+        started_at: child.startedAt,
+      })),
+      tokens_used: tokensUsed,
+      budget_breached: breached,
+    }),
+  }))
+
+  ctx.tools.register(defineTool({
+    name: 'stop_agent',
+    description: 'Gracefully cancel a child whose work is redundant or misdirected. Prefer send_message_to_agent to redirect a child that is merely off-track.',
+    parameters: {
+      agent_id: { type: 'string', required: true, description: 'Child id to cancel.' },
+    },
+    output: {
+      schema: {
+        type: 'object',
+        properties: {
+          success: { type: 'boolean', required: true },
+          error: { type: 'string' },
+        },
+        additionalProperties: false,
+      },
+      render: (_args, value) => {
+        const result = value as { success: boolean; error?: string }
+        return [{ type: 'text', text: result.success ? 'child cancelled' : `stop_agent failed: ${result.error ?? 'unknown'}` }]
+      },
+    },
+    execute: async (args, exec) => {
+      const child = children.get(args.agent_id)
+      if (child === undefined) return { success: false, error: `No child agent '${args.agent_id}'.` }
+      if (child.status !== 'running') return { success: false, error: `Child '${args.agent_id}' is already ${child.status}.` }
+      try {
+        const parentSessionId = exec.agent?.session.id
+        subagents().interrupt(args.agent_id as never, { kind: 'user', parentSessionId })
+        child.status = 'stopped'
+        return { success: true }
+      } catch (error) {
+        return { success: false, error: String(error instanceof Error ? error.message : error) }
+      }
+    },
+  }))
+
+  ctx.tools.register(defineTool({
+    name: 'agent_finish',
+    description: 'Child agents: submit your structured completion summary and end your turn. The final message IS the completion report the parent receives — this tool records that you are done.',
+    parameters: {
+      summary: { type: 'string', required: true, description: 'The structured completion report (findings, coverage, open items).' },
+    },
+    output: {
+      schema: {
+        type: 'object',
+        properties: {
+          success: { type: 'boolean', required: true },
+          message: { type: 'string' },
+        },
+        additionalProperties: false,
+      },
+      render: (_args, value) => {
+        const result = value as { message?: string }
+        return [{ type: 'text', text: String(result.message ?? 'completion recorded') }]
+      },
+    },
+    execute: async (args, exec) => {
+      // agent_finish executes inside the CHILD agent; its session id is the
+      // create_agent run id, so the summary lands on the tracked roster row.
+      const sessionKey = exec.agent === undefined ? undefined : String(exec.agent.session.id)
+      const child = sessionKey === undefined ? undefined : children.get(sessionKey)
+      if (child !== undefined && child.completionReport === undefined) child.completionReport = args.summary
+      return { success: true, message: `Completion recorded (${String(args.summary.length)} chars). End your turn now — your final message is the report.` }
+    },
+  }))
+
+  return handle
+}
+