Skip to content

feat(permissions): optimize auto-mode classifier with in-CWD fast path and two-stage classification #1098

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
Zilv1nas wants to merge 1 commit into
base: master
Choose a base branch
from
Draft

feat(permissions): optimize auto-mode classifier with in-CWD fast path and two-stage classification #1098

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
156 changes: 130 additions & 26 deletions src/extensions/permissions/classifier.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ describe("classifyToolCall", () => {
vi.useRealTimers()
})

it("returns safe verdict on first attempt", async () => {
it("returns safe verdict on Stage 1 without calling Stage 2", async () => {
completeMock.mockResolvedValue(
fakeResponse({ stopReason: "stop", content: '{"verdict":"safe","riskScore":"low","reason":"fine"}' }),
)
Expand All @@ -65,9 +65,98 @@ describe("classifyToolCall", () => {
expect(result.verdict).toBe("safe")
expect(result.riskScore).toBe("low")
expect(result.ok).toBe(true)
expect(result.stage).toBe(1)
expect(completeMock).toHaveBeenCalledTimes(1)
})

it("Stage 1 safe skips Stage 2 entirely", async () => {
completeMock.mockResolvedValueOnce(
fakeResponse({ stopReason: "stop", content: '{"verdict":"safe","reason":"stage1"}' }),
)

const result = await classifyToolCall(
fakeRegistry(),
{ toolName: "bash", input: { command: "ls" }, cwd: "/tmp" },
{ timeoutMs: 5000 },
)

expect(result.verdict).toBe("safe")
expect(result.stage).toBe(1)
expect(completeMock).toHaveBeenCalledTimes(1)
})

it("Stage 1 requires-confirmation falls through to Stage 2", async () => {
completeMock
.mockResolvedValueOnce(
fakeResponse({ stopReason: "stop", content: '{"verdict":"requires-confirmation","reason":"stage1 unsure"}' }),
)
.mockResolvedValueOnce(
fakeResponse({ stopReason: "stop", content: '{"verdict":"safe","riskScore":"low","reason":"stage2 safe"}' }),
)

const result = await classifyToolCall(
fakeRegistry(),
{ toolName: "bash", input: { command: "ls" }, cwd: "/tmp" },
{ timeoutMs: 5000 },
)

expect(result.verdict).toBe("safe")
expect(result.stage).toBe(2)
expect(completeMock).toHaveBeenCalledTimes(2)
})

it("Stage 1 parse failure falls through to Stage 2", async () => {
completeMock
.mockResolvedValueOnce(fakeResponse({ stopReason: "stop", content: "garbage text" }))
.mockResolvedValueOnce(
fakeResponse({ stopReason: "stop", content: '{"verdict":"safe","riskScore":"low","reason":"ok"}' }),
)

const result = await classifyToolCall(
fakeRegistry(),
{ toolName: "bash", input: { command: "ls" }, cwd: "/tmp" },
{ timeoutMs: 5000 },
)

expect(result.verdict).toBe("safe")
expect(result.stage).toBe(2)
expect(completeMock).toHaveBeenCalledTimes(2)
})

it("Stage 1 error falls through to Stage 2", async () => {
completeMock
.mockResolvedValueOnce(fakeResponse({ stopReason: "error", errorMessage: "stage1 error" }))
.mockResolvedValueOnce(
fakeResponse({ stopReason: "stop", content: '{"verdict":"safe","riskScore":"low","reason":"ok"}' }),
)

const result = await classifyToolCall(
fakeRegistry(),
{ toolName: "bash", input: { command: "ls" }, cwd: "/tmp" },
{ timeoutMs: 5000 },
)

expect(result.verdict).toBe("safe")
expect(result.stage).toBe(2)
expect(completeMock).toHaveBeenCalledTimes(2)
})

it("Stage 1 uses maxTokens=64 in request options", async () => {
let capturedOptions: unknown
completeMock.mockImplementation((_model: unknown, _context: unknown, options: unknown) => {
capturedOptions = options
return fakeResponse({ stopReason: "stop", content: '{"verdict":"safe","reason":"fine"}' })
})

await classifyToolCall(
fakeRegistry(),
{ toolName: "bash", input: { command: "ls" }, cwd: "/tmp" },
{ timeoutMs: 5000 },
)

expect((capturedOptions as { maxTokens?: number }).maxTokens).toBe(64)
})

it("keeps classifier tags while omitting Pi token limits for Kimchi", async () => {
let sentPayload: unknown
completeMock.mockImplementation((_model: unknown, _context: unknown, options: unknown) => {
Expand All @@ -85,7 +174,9 @@ describe("classifyToolCall", () => {
expect(sentPayload).toEqual({ tags: ["source:classifier", "existing"] })
})

it("retries up to 3 times on abort before giving up", async () => {
it("retries up to 3 times on abort before giving up (Stage 2)", async () => {
// Stage 1 fails (aborted → non-fatal fallthrough)
// Stage 2 retries 3 times, all aborted
completeMock.mockResolvedValue(fakeResponse({ stopReason: "aborted" }))

const promise = classifyToolCall(
Expand All @@ -101,15 +192,18 @@ describe("classifyToolCall", () => {
expect(result.ok).toBe(false)
expect(result.reason).toContain("classifier timeout")
expect(result.reason).toContain(CLASSIFIER_PRIMARY_MODEL_ID)
expect(completeMock).toHaveBeenCalledTimes(3)
// Stage1(1) + Stage2 retries(3) = 4
expect(completeMock).toHaveBeenCalledTimes(4)
})

it("succeeds on 2nd attempt after first abort", async () => {
it("succeeds on 2nd Stage 2 attempt after first abort", async () => {
// Stage 1 aborted (non-fatal) + Stage 2 first attempt aborted + Stage 2 second succeeds
completeMock
.mockResolvedValueOnce(fakeResponse({ stopReason: "aborted" }))
.mockResolvedValueOnce(fakeResponse({ stopReason: "aborted" })) // Stage 1
.mockResolvedValueOnce(fakeResponse({ stopReason: "aborted" })) // Stage 2 attempt 1
.mockResolvedValueOnce(
fakeResponse({ stopReason: "stop", content: '{"verdict":"safe","riskScore":"low","reason":"fine"}' }),
)
) // Stage 2 attempt 2

const promise = classifyToolCall(
fakeRegistry(),
Expand All @@ -123,16 +217,18 @@ describe("classifyToolCall", () => {
expect(result.verdict).toBe("safe")
expect(result.riskScore).toBe("low")
expect(result.ok).toBe(true)
expect(completeMock).toHaveBeenCalledTimes(2)
expect(result.stage).toBe(2)
expect(completeMock).toHaveBeenCalledTimes(3)
})

it("succeeds on 3rd attempt after two aborts", async () => {
it("succeeds on 3rd Stage 2 attempt after two aborts", async () => {
completeMock
.mockResolvedValueOnce(fakeResponse({ stopReason: "aborted" }))
.mockResolvedValueOnce(fakeResponse({ stopReason: "aborted" }))
.mockResolvedValueOnce(fakeResponse({ stopReason: "aborted" })) // Stage 1
.mockResolvedValueOnce(fakeResponse({ stopReason: "aborted" })) // Stage 2 attempt 1
.mockResolvedValueOnce(fakeResponse({ stopReason: "aborted" })) // Stage 2 attempt 2
.mockResolvedValueOnce(
fakeResponse({ stopReason: "stop", content: '{"verdict":"safe","riskScore":"low","reason":"fine"}' }),
)
) // Stage 2 attempt 3

const promise = classifyToolCall(
fakeRegistry(),
Expand All @@ -146,17 +242,19 @@ describe("classifyToolCall", () => {
expect(result.verdict).toBe("safe")
expect(result.riskScore).toBe("low")
expect(result.ok).toBe(true)
expect(completeMock).toHaveBeenCalledTimes(3)
expect(result.stage).toBe(2)
expect(completeMock).toHaveBeenCalledTimes(4)
})

it("calls fallback model after 3 retryable failures and returns its result", async () => {
it("calls fallback model after 3 Stage 2 retryable failures", async () => {
completeMock
.mockResolvedValueOnce(fakeResponse({ stopReason: "aborted" }))
.mockResolvedValueOnce(fakeResponse({ stopReason: "aborted" }))
.mockResolvedValueOnce(fakeResponse({ stopReason: "aborted" }))
.mockResolvedValueOnce(fakeResponse({ stopReason: "aborted" })) // Stage 1
.mockResolvedValueOnce(fakeResponse({ stopReason: "aborted" })) // Stage 2 attempt 1
.mockResolvedValueOnce(fakeResponse({ stopReason: "aborted" })) // Stage 2 attempt 2
.mockResolvedValueOnce(fakeResponse({ stopReason: "aborted" })) // Stage 2 attempt 3
.mockResolvedValueOnce(
fakeResponse({ stopReason: "stop", content: '{"verdict":"safe","riskScore":"low","reason":"fine"}' }),
)
) // Fallback

const promise = classifyToolCall(
fakeRegistry([fakeModel(CLASSIFIER_PRIMARY_MODEL_ID), fakeModel(CLASSIFIER_FALLBACK_MODEL_ID)]),
Expand All @@ -170,7 +268,7 @@ describe("classifyToolCall", () => {
expect(result.verdict).toBe("safe")
expect(result.riskScore).toBe("low")
expect(result.ok).toBe(true)
expect(completeMock).toHaveBeenCalledTimes(4)
expect(completeMock).toHaveBeenCalledTimes(5)
})

it("returns last result when no fallback model is provided", async () => {
Expand All @@ -187,7 +285,8 @@ describe("classifyToolCall", () => {

expect(result.verdict).toBe("requires-confirmation")
expect(result.ok).toBe(false)
expect(completeMock).toHaveBeenCalledTimes(3)
// Stage1(1) + Stage2 retries(3) = 4
expect(completeMock).toHaveBeenCalledTimes(4)
})

it("returns 'classifier aborted' when signal aborts during final attempt", async () => {
Expand All @@ -198,8 +297,8 @@ describe("classifyToolCall", () => {
let callCount = 0
completeMock.mockImplementation(() => {
callCount++
if (callCount === 3) {
// Simulate the outer signal aborting during the final classifier call
if (callCount === 4) {
// Simulate the outer signal aborting during the final Stage 2 attempt
controller.abort()
}
return fakeResponse({ stopReason: "aborted" })
Expand All @@ -215,7 +314,8 @@ describe("classifyToolCall", () => {
await vi.runAllTimersAsync()
const result = await promise

expect(completeMock).toHaveBeenCalledTimes(3)
// Stage1(1) + Stage2 retries(3) = 4
expect(completeMock).toHaveBeenCalledTimes(4)
expect(result.verdict).toBe("requires-confirmation")
expect(result.ok).toBe(false)
expect(result.reason).toBe("classifier aborted")
Expand All @@ -236,22 +336,26 @@ describe("classifyToolCall", () => {
expect(result.reason).toBe("classifier aborted")
})

it("does not retry on error and returns requires-confirmation", async () => {
completeMock.mockResolvedValue(fakeResponse({ stopReason: "error", errorMessage: "rate limit exceeded" }))
it("does not retry on error and returns requires-confirmation (Stage 2)", async () => {
// Stage 1 fails non-fatally, Stage 2 returns error
completeMock
.mockResolvedValueOnce(fakeResponse({ stopReason: "error", errorMessage: "stage1 error" }))
.mockResolvedValue(fakeResponse({ stopReason: "error", errorMessage: "rate limit exceeded" }))

const result = await classifyToolCall(
fakeRegistry(),
{ toolName: "bash", input: { command: "ls" }, cwd: "/tmp" },
{ timeoutMs: 5000 },
)

expect(completeMock).toHaveBeenCalledTimes(1)
expect(completeMock).toHaveBeenCalledTimes(2)
expect(result.verdict).toBe("requires-confirmation")
expect(result.ok).toBe(false)
expect(result.reason).toContain("classifier error: rate limit exceeded")
})

it("still falls back to unparseable when text is garbage", async () => {
it("still falls back to unparseable when text is garbage (Stage 2)", async () => {
// Stage 1 garbage (non-fatal), Stage 2 also garbage
completeMock.mockResolvedValue(fakeResponse({ stopReason: "stop", content: "not json at all" }))

const result = await classifyToolCall(
Expand Down
105 changes: 104 additions & 1 deletion src/extensions/permissions/classifier.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,13 @@ export const CLASSIFIER_REQUEST_TAG = "source:classifier"
export const CLASSIFIER_PRIMARY_MODEL_ID = "deepseek-v4-flash"
export const CLASSIFIER_FALLBACK_MODEL_ID = "minimax-m3"

/** Max tokens for Stage 1 (fast) classifier — just enough for a JSON verdict. */
const STAGE1_MAX_TOKENS = 64

/** Prompt suffix appended to the system prompt for Stage 1 (fast) classification. */
const STAGE1_PROMPT_SUFFIX =
'\n\nRespond with ONLY a JSON object: {"verdict":"safe"} or {"verdict":"requires-confirmation"}. No reasoning needed. Be conservative — if unsure, return requires-confirmation.'

export interface ClassifyInput {
toolName: string
input: Record<string, unknown>
Expand Down Expand Up @@ -40,6 +47,17 @@ export async function classifyToolCall(

if (signal?.aborted) return unavailable("classifier aborted")

// Stage 1 (fast): lightweight classifier call with minimal output.
// If it returns "safe", skip Stage 2 entirely — no second GPU call needed.
if (signal?.aborted) return unavailable("classifier aborted")
const stage1Result = await runClassifierFast(primaryModel, auth, call, options, signal)
if (stage1Result.ok && stage1Result.verdict === "safe") {
return { ...stage1Result, stage: 1 }
}

// Stage 2 (full reasoning): the existing full classifier call with retries

@kimchi-review kimchi-review Bot Aug 27, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ℹ️🔧 Maintainability

The if (signal?.aborted) return unavailable("classifier aborted") check is duplicated immediately before the Stage 1 call. The first check at the top of classifyToolCall already returns in the same way, so the second check is dead code that can confuse readers.

💡 Suggestion: Remove the redundant signal?.aborted check on line 58, leaving only the initial guard.

// and fallback. Used when Stage 1 does not return safe (either
// requires-confirmation or parse failure).
const maxAttempts = 3
let lastResult: InternalResult = unavailable("classifier unavailable")

Expand All @@ -50,7 +68,7 @@ export async function classifyToolCall(
}

const result = await runClassifier(primaryModel, auth, call, options, signal)
if (result.ok) return result
if (result.ok) return { ...result, stage: 2 }

if (!result.retryable) return result

Expand All @@ -69,6 +87,91 @@ export async function classifyToolCall(
return lastResult
}

/**
* Stage 1 (fast): lightweight classifier call with minimal prompt suffix and
* low max_tokens. Returns immediately if the verdict is "safe". Falls through
* to Stage 2 on any other outcome or failure.
*/
async function runClassifierFast(
model: Model<Api>,
auth: Awaited<ReturnType<ModelRegistry["getApiKeyAndHeaders"]>>,
call: ClassifyInput,
options: ClassifierOptions,
signal?: AbortSignal,
): Promise<InternalResult> {
if (!auth.ok || !auth.apiKey) return unavailable("no API key for classifier")
if (signal?.aborted) return unavailable("classifier aborted")

const controller = new AbortController()
const timeoutHandle = setTimeout(() => controller.abort(), options.timeoutMs)
const onOuterAbort = () => controller.abort()
signal?.addEventListener("abort", onOuterAbort)

try {
const response = await complete(
model,
{
systemPrompt: classifierSystemPrompt + STAGE1_PROMPT_SUFFIX,
messages: [
{
role: "user",
content: [{ type: "text", text: buildUserPrompt(call) }],
timestamp: Date.now(),
},
],
},
{
apiKey: auth.apiKey,
headers: auth.headers,
signal: controller.signal,
maxTokens: STAGE1_MAX_TOKENS,
onPayload: (payload: unknown) => {
if (payload && typeof payload === "object") {
const p = payload as Record<string, unknown>
const existing = Array.isArray(p.tags) ? (p.tags as string[]) : []
p.tags = [CLASSIFIER_REQUEST_TAG, ...existing]
}
return omitKimchiMaxTokensFromPayload(payload, model.provider)
},
},
)

if (response.stopReason === "aborted" || response.stopReason === "error") {
// Stage 1 failures are non-fatal — fall through to Stage 2.
return {
verdict: "requires-confirmation",
reason: "stage 1 failed, falling through",
ok: false,
retryable: false,
}
}

const text = response.content
.filter((c): c is { type: "text"; text: string } => c.type === "text")
.map((c) => c.text)
.join("\n")

const result = parseClassifierOutput(text)
if (result.ok && result.verdict === "safe") {
return { ...result, retryable: false }
}

// Stage 1 returned requires-confirmation or failed to parse — fall through.
return {
verdict: "requires-confirmation",
reason: "stage 1 inconclusive, falling through",
ok: false,
retryable: false,
}
} catch (_err) {
// Stage 1 errors are non-fatal — fall through to Stage 2.
return { verdict: "requires-confirmation", reason: "stage 1 error, falling through", ok: false, retryable: false }
} finally {
clearTimeout(timeoutHandle)

@kimchi-review kimchi-review Bot Aug 27, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ℹ️⚠️ Error Handling

runClassifierFast swallows Stage 1 exceptions and error/aborted stop reasons without logging the underlying cause. If Stage 1 starts consistently falling through in production, there is no diagnostic signal to explain why the optimization is not short-circuiting.

💡 Suggestion: Emit a debug or trace log in the catch block and in the error/aborted branches that includes the error message or stop reason, while still returning the non-fatal fallthrough result.

signal?.removeEventListener("abort", onOuterAbort)
}
}

async function runClassifier(
model: Model<Api>,
auth: Awaited<ReturnType<ModelRegistry["getApiKeyAndHeaders"]>>,
Expand Down
Loading
Loading