Skip to content

Commit e95a0dd

Browse files
jahoomajahooma
committed
Update checks for free mode cost
1 parent e358331 commit e95a0dd

File tree

2 files changed

+138
-28
lines changed

2 files changed

+138
-28
lines changed

‎common/src/constants/free-agents.ts‎

Lines changed: 116 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -7,30 +7,51 @@ import type { CostMode } from './model-config'
77
*/
88
export const FREE_COST_MODE = 'free' as const
99

10+
/**
11+
* Agents that are allowed to run in FREE mode.
12+
* Only these specific agents (and their expected models) get 0 credits in FREE mode.
13+
* This prevents abuse by users trying to use arbitrary agents for free.
14+
*
15+
* The mapping also specifies which models each agent is allowed to use in free mode.
16+
* If an agent uses a different model, it will be charged full credits.
17+
*/
18+
export const FREE_MODE_AGENT_MODELS: Record<string, Set<string>> = {
19+
// Root orchestrator
20+
'base2-free': new Set(['x-ai/grok-4.1-fast']),
21+
22+
// File exploration agents
23+
'file-picker': new Set(['google/gemini-2.5-flash-lite']),
24+
'file-picker-max': new Set(['x-ai/grok-4.1-fast']),
25+
'file-lister': new Set(['x-ai/grok-4.1-fast']),
26+
27+
// Research agents
28+
'researcher-web': new Set(['x-ai/grok-4.1-fast']),
29+
'researcher-docs': new Set(['x-ai/grok-4.1-fast']),
30+
31+
// Command execution
32+
'commander-lite': new Set(['x-ai/grok-4.1-fast']),
33+
34+
// Editor for free mode
35+
'editor-glm': new Set(['z-ai/glm-4.7', 'z-ai/glm-4.6']),
36+
}
37+
38+
/**
39+
* Set of all agent IDs allowed in FREE mode.
40+
* Derived from FREE_MODE_AGENT_MODELS for quick lookups.
41+
*/
42+
export const FREE_MODE_ALLOWED_AGENTS = new Set(Object.keys(FREE_MODE_AGENT_MODELS))
43+
1044
/**
1145
* Models that are allowed in FREE mode.
12-
* Only these cheap/fast models get 0 credits in FREE mode.
46+
* Derived from FREE_MODE_AGENT_MODELS - this is the union of all allowed models.
1347
* This prevents abuse by users trying to use expensive models for free.
1448
*/
15-
export const FREE_MODE_ALLOWED_MODELS = new Set([
16-
// Grok models used by base2-free, commander-lite, file-lister, file-picker-max
17-
'x-ai/grok-4.1-fast',
18-
'x-ai/grok-4-fast', // researcher agents
19-
20-
// Gemini flash models used by file-picker and other subagents
21-
'google/gemini-2.5-flash',
22-
'google/gemini-2.5-flash-lite',
23-
'google/gemini-2.5-flash-preview-09-2025',
24-
'google/gemini-2.5-flash-lite-preview-09-2025',
25-
26-
// GPT models used by editor-gpt-5, thinker, context-pruner
27-
'openai/gpt-5.1',
28-
'openai/gpt-5.1-chat',
29-
'openai/gpt-5-mini',
30-
])
49+
export const FREE_MODE_ALLOWED_MODELS = new Set(
50+
Object.values(FREE_MODE_AGENT_MODELS).flatMap((models) => Array.from(models)),
51+
)
3152

3253
/**
33-
* Agents that don't charge credits.
54+
* Agents that don't charge credits when credits would be very small (<5).
3455
*
3556
* These are typically lightweight utility agents that:
3657
* - Use cheap models (e.g., Gemini Flash)
@@ -39,6 +60,10 @@ export const FREE_MODE_ALLOWED_MODELS = new Set([
3960
*
4061
* Making them free avoids user confusion when they connect their own
4162
* Claude subscription (BYOK) but still see credit charges for non-Claude models.
63+
*
64+
* NOTE: This is separate from FREE_MODE_ALLOWED_AGENTS which is for the
65+
* explicit "free" cost mode. These agents get free credits only when
66+
* the cost would be trivial (<5 credits).
4267
*/
4368
export const FREE_TIER_AGENTS = new Set([
4469
'file-picker',
@@ -65,13 +90,83 @@ export function isFreeModeAllowedModel(model: string): boolean {
6590
}
6691

6792
/**
68-
* Check if an agent should be free (no credit charge).
93+
* Check if an agent is allowed to run in FREE mode.
94+
* Validates both the agent ID and optionally the publisher.
95+
*
96+
* For security, we only allow:
97+
* - Internal agents (no publisher, e.g., 'base2-free')
98+
* - Codebuff-published agents (publisher === 'codebuff')
99+
*
100+
* This prevents attackers from creating agents with matching names
101+
* under different publishers to abuse free mode.
102+
*/
103+
export function isFreeModeAllowedAgent(fullAgentId: string): boolean {
104+
const { publisherId, agentId } = parseAgentId(fullAgentId)
105+
106+
// Must have a valid agent ID
107+
if (!agentId) return false
108+
109+
// Must be in the allowed agents list
110+
if (!FREE_MODE_ALLOWED_AGENTS.has(agentId)) return false
111+
112+
// Must be either internal (no publisher) or from codebuff
113+
if (publisherId && publisherId !== 'codebuff') return false
114+
115+
return true
116+
}
117+
118+
/**
119+
* Check if a specific agent is allowed to use a specific model in FREE mode.
120+
* This is the strictest check - validates both the agent AND model combination.
121+
*
122+
* Returns true only if:
123+
* 1. The agent is allowed in free mode (isFreeModeAllowedAgent)
124+
* 2. The model is in that agent's allowed model set
125+
*/
126+
export function isFreeModeAllowedAgentModel(
127+
fullAgentId: string,
128+
model: string,
129+
): boolean {
130+
// First check if agent is allowed in free mode (includes publisher validation)
131+
if (!isFreeModeAllowedAgent(fullAgentId)) return false
132+
133+
// Parse to get the base agent ID for model lookup
134+
const { agentId } = parseAgentId(fullAgentId)
135+
if (!agentId) return false
136+
137+
// Get the allowed models for this agent
138+
const allowedModels = FREE_MODE_AGENT_MODELS[agentId]
139+
if (!allowedModels) return false
140+
141+
// Empty set means programmatic agent (no LLM calls expected)
142+
// For these, any model check should fail (they shouldn't be making LLM calls)
143+
if (allowedModels.size === 0) return false
144+
145+
return allowedModels.has(model)
146+
}
147+
148+
/**
149+
* Check if an agent should be free (no credit charge) for small requests.
150+
* This is separate from FREE mode - these agents get free credits only
151+
* when the cost would be trivial (<5 credits).
152+
*
69153
* Handles all agent ID formats:
70154
* - 'file-picker'
71155
* - 'file-picker@1.0.0'
72156
* - 'codebuff/file-picker@0.0.2'
73157
*/
74158
export function isFreeAgent(fullAgentId: string): boolean {
75-
const { agentId } = parseAgentId(fullAgentId)
76-
return agentId ? FREE_TIER_AGENTS.has(agentId) : false
159+
const { publisherId, agentId } = parseAgentId(fullAgentId)
160+
161+
// Must have a valid agent ID
162+
if (!agentId) return false
163+
164+
// Must be in the free tier agents list
165+
if (!FREE_TIER_AGENTS.has(agentId)) return false
166+
167+
// Must be either internal (no publisher) or from codebuff
168+
// This prevents publisher spoofing attacks
169+
if (publisherId && publisherId !== 'codebuff') return false
170+
171+
return true
77172
}

‎web/src/llm-api/helpers.ts‎

Lines changed: 22 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,10 @@
11
import { setupBigQuery } from '@codebuff/bigquery'
22
import { consumeCreditsAndAddAgentStep } from '@codebuff/billing'
3-
import { isFreeAgent, isFreeMode, isFreeModeAllowedModel } from '@codebuff/common/constants/free-agents'
3+
import {
4+
isFreeAgent,
5+
isFreeMode,
6+
isFreeModeAllowedAgentModel,
7+
} from '@codebuff/common/constants/free-agents'
48
import { PROFIT_MARGIN } from '@codebuff/common/old-constants'
59

610
import type { InsertMessageBigqueryFn } from '@codebuff/common/types/contracts/bigquery'
@@ -123,13 +127,24 @@ export async function consumeCreditsForMessage(params: {
123127
costMode,
124128
} = params
125129

126-
// FREE mode: agents using allowed models cost 0 credits
127-
// Only whitelisted cheap models (grok-4.1-fast, gemini flash, gpt-5.1, etc.) are free
128-
// This prevents abuse by using expensive models in FREE mode
129-
// Free tier agents (like file-picker) also don't charge credits
130+
// Calculate initial credits based on cost
130131
const initialCredits = Math.round(usageData.cost * 100 * (1 + PROFIT_MARGIN))
131-
const isFreeModeAndAllowed = isFreeMode(costMode) && isFreeModeAllowedModel(model)
132-
const credits = isFreeModeAndAllowed || (isFreeAgent(agentId) && initialCredits < 5) ? 0 : initialCredits
132+
133+
// FREE mode: only specific agents using their expected models cost 0 credits
134+
// This is the strictest check - validates:
135+
// 1. The cost mode is 'free'
136+
// 2. The agent is in the allowed free-mode agents list
137+
// 3. The model matches what that specific agent is allowed to use
138+
// 4. The agent is either internal or published by 'codebuff' (prevents publisher spoofing)
139+
const isFreeModeAndAllowed =
140+
isFreeMode(costMode) && isFreeModeAllowedAgentModel(agentId, model)
141+
142+
// Free tier agents (like file-picker) also don't charge credits for small requests
143+
// This is separate from FREE mode and helps with BYOK users
144+
// Also validates publisher to prevent spoofing attacks
145+
const isFreeAgentSmallRequest = isFreeAgent(agentId) && initialCredits < 5
146+
147+
const credits = isFreeModeAndAllowed || isFreeAgentSmallRequest ? 0 : initialCredits
133148

134149
await consumeCreditsAndAddAgentStep({
135150
messageId,

0 commit comments

Comments
 (0)