diff --git a/cli/src/components/waiting-room-screen.tsx b/cli/src/components/waiting-room-screen.tsx index 32477a798..7cc0aca4a 100644 --- a/cli/src/components/waiting-room-screen.tsx +++ b/cli/src/components/waiting-room-screen.tsx @@ -17,6 +17,7 @@ import { exitFreebuffCleanly } from '../utils/freebuff-exit' import { getLogoAccentColor, getLogoBlockColor } from '../utils/theme-system' import type { FreebuffSessionResponse } from '../types/freebuff-session' +import type { FreebuffIpPrivacySignal } from '@codebuff/common/types/freebuff-session' interface WaitingRoomScreenProps { session: FreebuffSessionResponse | null @@ -55,6 +56,35 @@ const formatRetryAfter = (ms: number): string => { return rem === 0 ? `${hours}h` : `${hours}h ${rem}m` } +const PRIVACY_SIGNAL_LABELS: Partial> = + { + anonymous: 'anonymized network', + proxy: 'proxy', + relay: 'relay', + res_proxy: 'residential proxy', + tor: 'Tor', + vpn: 'VPN', + } + +const formatPrivacySignalList = ( + signals: FreebuffIpPrivacySignal[] | undefined, +): string => { + const labels = Array.from( + new Set( + signals + ?.map((signal) => PRIVACY_SIGNAL_LABELS[signal]) + .filter((label): label is string => Boolean(label)) ?? [], + ), + ) + + if (labels.length === 0) { + return 'VPN, Tor, proxy, relay, or anonymized network' + } + if (labels.length === 1) return labels[0] + if (labels.length === 2) return `${labels[0]} or ${labels[1]}` + return `${labels.slice(0, -1).join(', ')}, or ${labels[labels.length - 1]}` +} + export const WaitingRoomScreen: React.FC = ({ session, error, @@ -263,7 +293,23 @@ export const WaitingRoomScreen: React.FC = ({ ⚠ Free mode isn't available in your region - {session.countryCode === 'UNKNOWN' ? ( + {session.countryBlockReason === 'anonymous_network' ? ( + <> + We detected{' '} + {formatPrivacySignalList(session.ipPrivacySignals)} traffic + {session.countryCode === 'UNKNOWN' ? ( + '' + ) : ( + <> + {' '} + from{' '} + {session.countryCode} + + )} + . Freebuff can't be used from anonymized networks. Press + Ctrl+C to exit. + + ) : session.countryCode === 'UNKNOWN' ? ( <> We couldn't verify an eligible location for this request. VPN, Tor, proxy, or unknown-location traffic can't use diff --git a/cli/src/hooks/helpers/send-message.ts b/cli/src/hooks/helpers/send-message.ts index a86870fe5..cf9063166 100644 --- a/cli/src/hooks/helpers/send-message.ts +++ b/cli/src/hooks/helpers/send-message.ts @@ -12,7 +12,7 @@ import { IS_FREEBUFF } from '../../utils/constants' import { processBashContext } from '../../utils/bash-context-processor' import { markRunningAgentsAsCancelled } from '../../utils/block-operations' import { - getCountryCodeFromFreeModeError, + getCountryBlockFromFreeModeError, getFreebuffGateErrorKind, isOutOfCreditsError, isFreeModeUnavailableError, @@ -394,7 +394,9 @@ export const handleRunCompletion = (params: { updater.setError(FREE_MODE_UNAVAILABLE_MESSAGE) if (IS_FREEBUFF) { markFreebuffSessionCountryBlocked( - getCountryCodeFromFreeModeError(output) ?? 'UNKNOWN', + getCountryBlockFromFreeModeError(output) ?? { + countryCode: 'UNKNOWN', + }, ) } finalizeAfterError() @@ -494,7 +496,9 @@ export const handleRunError = (params: { updater.setError(FREE_MODE_UNAVAILABLE_MESSAGE) if (IS_FREEBUFF) { markFreebuffSessionCountryBlocked( - getCountryCodeFromFreeModeError(error) ?? 'UNKNOWN', + getCountryBlockFromFreeModeError(error) ?? { + countryCode: 'UNKNOWN', + }, ) } return diff --git a/cli/src/hooks/use-freebuff-session.ts b/cli/src/hooks/use-freebuff-session.ts index 119e769b8..463a49126 100644 --- a/cli/src/hooks/use-freebuff-session.ts +++ b/cli/src/hooks/use-freebuff-session.ts @@ -16,6 +16,10 @@ import { logger } from '../utils/logger' import { saveFreebuffModelPreference } from '../utils/settings' import type { FreebuffSessionResponse } from '../types/freebuff-session' +import type { + FreebuffCountryBlockReason, + FreebuffIpPrivacySignal, +} from '@codebuff/common/types/freebuff-session' const POLL_INTERVAL_QUEUED_MS = 5_000 const POLL_INTERVAL_ACTIVE_MS = 30_000 @@ -319,10 +323,14 @@ export function markFreebuffSessionSuperseded(): void { * Transitioning the session state here unmounts the Chat surface in favor of * the waiting-room's country_blocked message, so the user can't keep typing * and sending doomed requests. */ -export function markFreebuffSessionCountryBlocked(countryCode: string): void { +export function markFreebuffSessionCountryBlocked(params: { + countryCode: string + countryBlockReason?: FreebuffCountryBlockReason + ipPrivacySignals?: FreebuffIpPrivacySignal[] +}): void { if (!IS_FREEBUFF) return controller?.abort() - controller?.apply({ status: 'country_blocked', countryCode }) + controller?.apply({ status: 'country_blocked', ...params }) // Best-effort DELETE so we don't hold a waiting-room seat on a session the // server is already refusing to serve at chat time. releaseFreebuffSlot().catch(() => {}) diff --git a/cli/src/utils/__tests__/error-handling.test.ts b/cli/src/utils/__tests__/error-handling.test.ts index 00097730b..190009326 100644 --- a/cli/src/utils/__tests__/error-handling.test.ts +++ b/cli/src/utils/__tests__/error-handling.test.ts @@ -3,6 +3,7 @@ import { describe, test, expect } from 'bun:test' import { isOutOfCreditsError, isFreeModeUnavailableError, + getCountryBlockFromFreeModeError, OUT_OF_CREDITS_MESSAGE, FREE_MODE_UNAVAILABLE_MESSAGE, createErrorMessage, @@ -70,7 +71,11 @@ describe('error-handling', () => { describe('isFreeModeUnavailableError', () => { test('returns true for error with statusCode 403 and error free_mode_unavailable', () => { - const error = { statusCode: 403, error: 'free_mode_unavailable', message: 'Free mode is not available in your country.' } + const error = { + statusCode: 403, + error: 'free_mode_unavailable', + message: 'Free mode is not available in your country.', + } expect(isFreeModeUnavailableError(error)).toBe(true) }) @@ -80,12 +85,20 @@ describe('error-handling', () => { }) test('returns false for 403 with different error code', () => { - const error = { statusCode: 403, error: 'account_suspended', message: 'Suspended' } + const error = { + statusCode: 403, + error: 'account_suspended', + message: 'Suspended', + } expect(isFreeModeUnavailableError(error)).toBe(false) }) test('returns false for non-403 status with free_mode_unavailable error', () => { - const error = { statusCode: 400, error: 'free_mode_unavailable', message: 'Bad request' } + const error = { + statusCode: 400, + error: 'free_mode_unavailable', + message: 'Bad request', + } expect(isFreeModeUnavailableError(error)).toBe(false) }) @@ -102,9 +115,51 @@ describe('error-handling', () => { }) }) + describe('getCountryBlockFromFreeModeError', () => { + test('extracts country block details from free-mode unavailable errors', () => { + const error = { + statusCode: 403, + error: 'free_mode_unavailable', + countryCode: 'US', + countryBlockReason: 'anonymous_network', + ipPrivacySignals: ['vpn', 'hosting', 123], + } + + expect(getCountryBlockFromFreeModeError(error)).toEqual({ + countryCode: 'US', + countryBlockReason: 'anonymous_network', + ipPrivacySignals: ['vpn', 'hosting'], + }) + }) + + test('defaults missing country code to UNKNOWN', () => { + const error = { + statusCode: 403, + error: 'free_mode_unavailable', + } + + expect(getCountryBlockFromFreeModeError(error)).toEqual({ + countryCode: 'UNKNOWN', + countryBlockReason: undefined, + ipPrivacySignals: undefined, + }) + }) + + test('returns null for non-free-mode errors', () => { + expect( + getCountryBlockFromFreeModeError({ + statusCode: 403, + error: 'account_suspended', + }), + ).toBe(null) + }) + }) + describe('FREE_MODE_UNAVAILABLE_MESSAGE', () => { test('mentions unavailability in country', () => { - expect(FREE_MODE_UNAVAILABLE_MESSAGE.toLowerCase()).toContain('not available in your country') + expect(FREE_MODE_UNAVAILABLE_MESSAGE.toLowerCase()).toContain( + 'not available in your country', + ) }) }) diff --git a/cli/src/utils/error-handling.ts b/cli/src/utils/error-handling.ts index 9b624ea52..742c5a507 100644 --- a/cli/src/utils/error-handling.ts +++ b/cli/src/utils/error-handling.ts @@ -1,6 +1,10 @@ import { env } from '@codebuff/common/env' import type { ChatMessage } from '../types/chat' +import type { + FreebuffCountryBlockReason, + FreebuffIpPrivacySignal, +} from '@codebuff/common/types/freebuff-session' import { IS_FREEBUFF } from './constants' @@ -57,20 +61,38 @@ export const isFreeModeUnavailableError = (error: unknown): boolean => { return false } -/** - * Extract the detected countryCode off a free_mode_unavailable error, if the - * server included one. Used to populate the country_blocked screen after the - * chat-completions gate rejects a user whose session-level country check did - * not catch the request first. - */ -export const getCountryCodeFromFreeModeError = ( +export const getCountryBlockFromFreeModeError = ( error: unknown, -): string | null => { +): { + countryCode: string + countryBlockReason?: FreebuffCountryBlockReason + ipPrivacySignals?: FreebuffIpPrivacySignal[] +} | null => { if (!isFreeModeUnavailableError(error)) return null - const candidate = (error as { countryCode?: unknown }).countryCode - return typeof candidate === 'string' && candidate.length > 0 - ? candidate - : null + const errorDetails = error as { + countryCode?: unknown + countryBlockReason?: unknown + ipPrivacySignals?: unknown + } + const countryCode = + typeof errorDetails.countryCode === 'string' && + errorDetails.countryCode.length > 0 + ? errorDetails.countryCode + : 'UNKNOWN' + + return { + countryCode, + countryBlockReason: + typeof errorDetails.countryBlockReason === 'string' + ? (errorDetails.countryBlockReason as FreebuffCountryBlockReason) + : undefined, + ipPrivacySignals: Array.isArray(errorDetails.ipPrivacySignals) + ? errorDetails.ipPrivacySignals.filter( + (signal): signal is FreebuffIpPrivacySignal => + typeof signal === 'string', + ) + : undefined, + } } /** diff --git a/common/src/types/freebuff-session.ts b/common/src/types/freebuff-session.ts index c7322b665..eff5abff7 100644 --- a/common/src/types/freebuff-session.ts +++ b/common/src/types/freebuff-session.ts @@ -21,6 +21,23 @@ export interface FreebuffSessionRateLimit { recentCount: number } +export type FreebuffCountryBlockReason = + | 'country_not_allowed' + | 'anonymized_or_unknown_country' + | 'anonymous_network' + | 'missing_client_ip' + | 'unresolved_client_ip' + +export type FreebuffIpPrivacySignal = + | 'anonymous' + | 'vpn' + | 'proxy' + | 'tor' + | 'relay' + | 'res_proxy' + | 'hosting' + | 'service' + export type FreebuffSessionServerResponse = | { /** Waiting room is globally off; free-mode requests flow through @@ -106,6 +123,8 @@ export type FreebuffSessionServerResponse = * screen. `countryCode` is the resolved country, or UNKNOWN. */ status: 'country_blocked' countryCode: string + countryBlockReason?: FreebuffCountryBlockReason + ipPrivacySignals?: FreebuffIpPrivacySignal[] } | { /** User has an active session bound to a different model. Returned diff --git a/common/src/types/session-state.ts b/common/src/types/session-state.ts index 3896f8788..a116a5cde 100644 --- a/common/src/types/session-state.ts +++ b/common/src/types/session-state.ts @@ -69,6 +69,9 @@ export const AgentOutputSchema = z.discriminatedUnion('type', [ message: z.string(), statusCode: z.number().optional(), error: z.string().optional(), + countryCode: z.string().optional(), + countryBlockReason: z.string().optional(), + ipPrivacySignals: z.array(z.string()).optional(), }), ]) export type AgentOutput = z.infer diff --git a/common/src/util/error.ts b/common/src/util/error.ts index 1861e1d39..610ff3208 100644 --- a/common/src/util/error.ts +++ b/common/src/util/error.ts @@ -198,18 +198,56 @@ export function unwrapPromptResult(result: PromptResult): T { export function parseApiErrorResponseBody(responseBody: unknown): { errorCode?: string message?: string + countryCode?: string + countryBlockReason?: string + ipPrivacySignals?: string[] } { if (typeof responseBody !== 'string') return {} try { const parsed: unknown = JSON.parse(responseBody) if (!parsed || typeof parsed !== 'object') return {} - const result: { errorCode?: string; message?: string } = {} - if ('error' in parsed && typeof (parsed as { error: unknown }).error === 'string') { + const result: { + errorCode?: string + message?: string + countryCode?: string + countryBlockReason?: string + ipPrivacySignals?: string[] + } = {} + if ( + 'error' in parsed && + typeof (parsed as { error: unknown }).error === 'string' + ) { result.errorCode = (parsed as { error: string }).error } - if ('message' in parsed && typeof (parsed as { message: unknown }).message === 'string') { + if ( + 'message' in parsed && + typeof (parsed as { message: unknown }).message === 'string' + ) { result.message = (parsed as { message: string }).message } + if ( + 'countryCode' in parsed && + typeof (parsed as { countryCode: unknown }).countryCode === 'string' + ) { + result.countryCode = (parsed as { countryCode: string }).countryCode + } + if ( + 'countryBlockReason' in parsed && + typeof (parsed as { countryBlockReason: unknown }).countryBlockReason === + 'string' + ) { + result.countryBlockReason = ( + parsed as { countryBlockReason: string } + ).countryBlockReason + } + if ('ipPrivacySignals' in parsed) { + const signals = (parsed as { ipPrivacySignals: unknown }).ipPrivacySignals + if (Array.isArray(signals)) { + result.ipPrivacySignals = signals.filter( + (signal): signal is string => typeof signal === 'string', + ) + } + } return result } catch { return {} diff --git a/packages/agent-runtime/src/__tests__/loop-agent-steps.test.ts b/packages/agent-runtime/src/__tests__/loop-agent-steps.test.ts index 63ddf60d2..873079f51 100644 --- a/packages/agent-runtime/src/__tests__/loop-agent-steps.test.ts +++ b/packages/agent-runtime/src/__tests__/loop-agent-steps.test.ts @@ -955,6 +955,9 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () => responseBody: JSON.stringify({ error: 'free_mode_unavailable', message: 'Free mode is not available in your country.', + countryCode: 'US', + countryBlockReason: 'anonymous_network', + ipPrivacySignals: ['vpn', 'hosting'], }), isRetryable: false, }) @@ -976,6 +979,9 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () => expect(result.output.error).toBe('free_mode_unavailable') // Should propagate the status code expect(result.output.statusCode).toBe(403) + expect(result.output.countryCode).toBe('US') + expect(result.output.countryBlockReason).toBe('anonymous_network') + expect(result.output.ipPrivacySignals).toEqual(['vpn', 'hosting']) } }) diff --git a/packages/agent-runtime/src/run-agent-step.ts b/packages/agent-runtime/src/run-agent-step.ts index 4b8267033..38af3ae65 100644 --- a/packages/agent-runtime/src/run-agent-step.ts +++ b/packages/agent-runtime/src/run-agent-step.ts @@ -1097,11 +1097,21 @@ export async function loopAgentSteps( let errorMessage = '' let errorCode: string | undefined + let countryCode: string | undefined + let countryBlockReason: string | undefined + let ipPrivacySignals: string[] | undefined let hasServerMessage = false if (error instanceof APICallError) { errorMessage = `${error.message}` const parsed = parseApiErrorResponseBody(error.responseBody) if (parsed.errorCode) errorCode = parsed.errorCode + if (parsed.countryCode) countryCode = parsed.countryCode + if (parsed.countryBlockReason) { + countryBlockReason = parsed.countryBlockReason + } + if (parsed.ipPrivacySignals) { + ipPrivacySignals = parsed.ipPrivacySignals + } if (parsed.message) { errorMessage = parsed.message hasServerMessage = true @@ -1139,6 +1149,9 @@ export async function loopAgentSteps( message: hasServerMessage ? errorMessage : 'Agent run error: ' + errorMessage, ...(statusCode !== undefined && { statusCode }), ...(errorCode !== undefined && { error: errorCode }), + ...(countryCode !== undefined && { countryCode }), + ...(countryBlockReason !== undefined && { countryBlockReason }), + ...(ipPrivacySignals !== undefined && { ipPrivacySignals }), }, } } diff --git a/sdk/src/__tests__/run-cancellation.test.ts b/sdk/src/__tests__/run-cancellation.test.ts index 46c7ed4bc..2eba5d3e4 100644 --- a/sdk/src/__tests__/run-cancellation.test.ts +++ b/sdk/src/__tests__/run-cancellation.test.ts @@ -1,4 +1,3 @@ - import * as mainPromptModule from '@codebuff/agent-runtime/main-prompt' import { withSystemTags } from '@codebuff/agent-runtime/util/messages' import { getInitialSessionState } from '@codebuff/common/types/session-state' @@ -37,9 +36,11 @@ describe('Run Cancellation Handling', () => { spyOn(databaseModule, 'addAgentStep').mockResolvedValue('step-1') // Server session state already includes the user's message (as the server would normally do) - const serverSessionState = getInitialSessionState(getStubProjectFileContext()) + const serverSessionState = getInitialSessionState( + getStubProjectFileContext(), + ) serverSessionState.mainAgentState.messageHistory.push( - userMessage('Please fix the bug'), // Server added this + userMessage('Please fix the bug'), // Server added this assistantMessage('I will help you with that.'), ) @@ -82,10 +83,10 @@ describe('Run Cancellation Handling', () => { const messageHistory = result.sessionState!.mainAgentState.messageHistory const userMessages = messageHistory.filter((m) => m.role === 'user') - + // Should have exactly 1 user message, not 2 expect(userMessages.length).toBe(1) - + // Total messages should be 2 (user + assistant), not 3 expect(messageHistory.length).toBe(2) }) @@ -107,9 +108,11 @@ describe('Run Cancellation Handling', () => { const abortController = new AbortController() // Server session state already includes the user's message (server processed it) - const serverSessionState = getInitialSessionState(getStubProjectFileContext()) + const serverSessionState = getInitialSessionState( + getStubProjectFileContext(), + ) serverSessionState.mainAgentState.messageHistory.push( - userMessage('Please fix the bug'), // Server added the user's message + userMessage('Please fix the bug'), // Server added the user's message assistantMessage('I will help you with that.'), ) @@ -131,7 +134,11 @@ describe('Run Cancellation Handling', () => { // Simulate agent runtime adding interruption message on abort serverSessionState.mainAgentState.messageHistory.push( - userMessage(withSystemTags("User interrupted the response. The assistant's previous work has been preserved.")) + userMessage( + withSystemTags( + "User interrupted the response. The assistant's previous work has been preserved.", + ), + ), ) // Server still responds with its session state @@ -169,16 +176,19 @@ describe('Run Cancellation Handling', () => { // The user's message should NOT be duplicated const messageHistory = result.sessionState!.mainAgentState.messageHistory - + // Count user messages (excluding system interruption messages) const userPromptMessages = messageHistory.filter( - (m) => m.role === 'user' && - m.content.some((c: any) => c.type === 'text' && c.text.includes('fix the bug')) + (m) => + m.role === 'user' && + m.content.some( + (c: any) => c.type === 'text' && c.text.includes('fix the bug'), + ), ) - + // Should have exactly 1 user message with the prompt, not 2 expect(userPromptMessages.length).toBe(1) - + // Total messages should be: 1 user + 1 assistant (original) + 1 interruption = 3 // The server state already has the content; pendingAgentResponse is not duplicated. expect(messageHistory.length).toBe(3) @@ -199,11 +209,17 @@ describe('Run Cancellation Handling', () => { spyOn(databaseModule, 'addAgentStep').mockResolvedValue('step-1') // Simulate AI SDK's AI_APICallError with responseBody (what the server returns for free_mode_unavailable) - const apiError = new Error('Forbidden') as Error & { statusCode: number; responseBody: string } + const apiError = new Error('Forbidden') as Error & { + statusCode: number + responseBody: string + } apiError.statusCode = 403 apiError.responseBody = JSON.stringify({ error: 'free_mode_unavailable', message: 'Free mode is not available in your country.', + countryCode: 'US', + countryBlockReason: 'anonymous_network', + ipPrivacySignals: ['vpn', 'hosting'], }) spyOn(mainPromptModule, 'callMainPrompt').mockRejectedValue(apiError) @@ -218,12 +234,23 @@ describe('Run Cancellation Handling', () => { }) expect(result.output.type).toBe('error') - const output = result.output as { type: 'error'; message: string; statusCode?: number; error?: string } + const output = result.output as { + type: 'error' + message: string + statusCode?: number + error?: string + countryCode?: string + countryBlockReason?: string + ipPrivacySignals?: string[] + } // Should use the message from the response body, not the generic "Forbidden" expect(output.message).toBe('Free mode is not available in your country.') expect(output.statusCode).toBe(403) // Should propagate the error code so isFreeModeUnavailableError can match expect(output.error).toBe('free_mode_unavailable') + expect(output.countryCode).toBe('US') + expect(output.countryBlockReason).toBe('anonymous_network') + expect(output.ipPrivacySignals).toEqual(['vpn', 'hosting']) }) it('extracts error code from responseBody for account_suspended 403', async () => { @@ -240,7 +267,10 @@ describe('Run Cancellation Handling', () => { spyOn(databaseModule, 'finishAgentRun').mockResolvedValue(undefined) spyOn(databaseModule, 'addAgentStep').mockResolvedValue('step-1') - const apiError = new Error('Forbidden') as Error & { statusCode: number; responseBody: string } + const apiError = new Error('Forbidden') as Error & { + statusCode: number + responseBody: string + } apiError.statusCode = 403 apiError.responseBody = JSON.stringify({ error: 'account_suspended', @@ -258,8 +288,15 @@ describe('Run Cancellation Handling', () => { prompt: 'hello', }) - const output = result.output as { type: 'error'; message: string; statusCode?: number; error?: string } - expect(output.message).toBe('Your account has been suspended due to billing issues.') + const output = result.output as { + type: 'error' + message: string + statusCode?: number + error?: string + } + expect(output.message).toBe( + 'Your account has been suspended due to billing issues.', + ) expect(output.statusCode).toBe(403) expect(output.error).toBe('account_suspended') }) @@ -278,7 +315,10 @@ describe('Run Cancellation Handling', () => { spyOn(databaseModule, 'finishAgentRun').mockResolvedValue(undefined) spyOn(databaseModule, 'addAgentStep').mockResolvedValue('step-1') - const apiError = new Error('Forbidden') as Error & { statusCode: number; responseBody: string } + const apiError = new Error('Forbidden') as Error & { + statusCode: number + responseBody: string + } apiError.statusCode = 403 apiError.responseBody = 'not valid json' @@ -293,7 +333,12 @@ describe('Run Cancellation Handling', () => { prompt: 'hello', }) - const output = result.output as { type: 'error'; message: string; statusCode?: number; error?: string } + const output = result.output as { + type: 'error' + message: string + statusCode?: number + error?: string + } expect(output.message).toBe('Forbidden') expect(output.statusCode).toBe(403) expect(output.error).toBeUndefined() @@ -329,7 +374,9 @@ describe('Run Cancellation Handling', () => { // Should return an error output expect(result.output.type).toBe('error') - expect((result.output as { type: 'error'; message: string }).message).toBe('Network connection failed') + expect((result.output as { type: 'error'; message: string }).message).toBe( + 'Network connection failed', + ) // The user's message should be preserved in the session state expect(result.sessionState).toBeDefined() @@ -345,7 +392,9 @@ describe('Run Cancellation Handling', () => { expect(userPromptMessage).toBeDefined() // Verify the message content contains the original prompt - const textContent = userPromptMessage!.content.find((c: any) => c.type === 'text') as { type: 'text'; text: string } | undefined + const textContent = userPromptMessage!.content.find( + (c: any) => c.type === 'text', + ) as { type: 'text'; text: string } | undefined expect(textContent).toBeDefined() expect(textContent!.text).toContain('Please fix the bug in my code') }) @@ -365,11 +414,14 @@ describe('Run Cancellation Handling', () => { spyOn(databaseModule, 'addAgentStep').mockResolvedValue('step-1') const abortController = new AbortController() - const serverSessionState = getInitialSessionState(getStubProjectFileContext()) + const serverSessionState = getInitialSessionState( + getStubProjectFileContext(), + ) serverSessionState.mainAgentState.messageHistory.push( userMessage('User prompt'), ) - const originalHistoryLength = serverSessionState.mainAgentState.messageHistory.length + const originalHistoryLength = + serverSessionState.mainAgentState.messageHistory.length spyOn(mainPromptModule, 'callMainPrompt').mockImplementation( async (params: Parameters[0]) => { @@ -380,7 +432,11 @@ describe('Run Cancellation Handling', () => { // Simulate agent runtime adding interruption message on abort serverSessionState.mainAgentState.messageHistory.push( - userMessage(withSystemTags("User interrupted the response. The assistant's previous work has been preserved.")) + userMessage( + withSystemTags( + "User interrupted the response. The assistant's previous work has been preserved.", + ), + ), ) await sendAction({ @@ -423,7 +479,9 @@ describe('Run Cancellation Handling', () => { // The last message should be the interruption (user role), not an empty assistant message const lastMessage = messageHistory[messageHistory.length - 1] expect(lastMessage.role).toBe('user') - expect((lastMessage.content[0] as { type: 'text'; text: string }).text).toContain('User interrupted') + expect( + (lastMessage.content[0] as { type: 'text'; text: string }).text, + ).toContain('User interrupted') // Verify there's no empty assistant message before the interruption const secondToLastMessage = messageHistory[messageHistory.length - 2] @@ -518,7 +576,9 @@ describe('Run Cancellation Handling', () => { const abortController = new AbortController() // Create a session state with some existing message history to verify it's preserved - const serverSessionState = getInitialSessionState(getStubProjectFileContext()) + const serverSessionState = getInitialSessionState( + getStubProjectFileContext(), + ) serverSessionState.mainAgentState.messageHistory.push( userMessage('User prompt'), assistantMessage('I will help you with that.'), @@ -541,10 +601,13 @@ describe('Run Cancellation Handling', () => { role: 'tool', toolCallId: 'tool-1', toolName: 'read_files', - content: [{ type: 'json', value: [{ path: 'file.ts', content: 'const x = 1;' }] }], + content: [ + { type: 'json', value: [{ path: 'file.ts', content: 'const x = 1;' }] }, + ], }) - const originalHistoryLength = serverSessionState.mainAgentState.messageHistory.length + const originalHistoryLength = + serverSessionState.mainAgentState.messageHistory.length spyOn(mainPromptModule, 'callMainPrompt').mockImplementation( async (params: Parameters[0]) => { @@ -564,7 +627,11 @@ describe('Run Cancellation Handling', () => { // Simulate agent runtime adding interruption message on abort serverSessionState.mainAgentState.messageHistory.push( - userMessage(withSystemTags("User interrupted the response. The assistant's previous work has been preserved.")) + userMessage( + withSystemTags( + "User interrupted the response. The assistant's previous work has been preserved.", + ), + ), ) // Server still sends the prompt-response with the full session state @@ -615,7 +682,9 @@ describe('Run Cancellation Handling', () => { const toolCallMessage = messageHistory.find( (m) => m.role === 'assistant' && - m.content.some((c: any) => c.type === 'tool-call' && c.toolCallId === 'tool-1'), + m.content.some( + (c: any) => c.type === 'tool-call' && c.toolCallId === 'tool-1', + ), ) expect(toolCallMessage).toBeDefined() @@ -644,7 +713,9 @@ describe('Run Cancellation Handling', () => { spyOn(databaseModule, 'addAgentStep').mockResolvedValue('step-1') const abortController = new AbortController() - const serverSessionState = getInitialSessionState(getStubProjectFileContext()) + const serverSessionState = getInitialSessionState( + getStubProjectFileContext(), + ) spyOn(mainPromptModule, 'callMainPrompt').mockImplementation( async (params: Parameters[0]) => { @@ -655,7 +726,11 @@ describe('Run Cancellation Handling', () => { // Simulate agent runtime adding interruption message on abort serverSessionState.mainAgentState.messageHistory.push( - userMessage(withSystemTags("User interrupted the response. The assistant's previous work has been preserved.")) + userMessage( + withSystemTags( + "User interrupted the response. The assistant's previous work has been preserved.", + ), + ), ) await sendAction({ @@ -697,7 +772,9 @@ describe('Run Cancellation Handling', () => { expect(lastMessage.role).toBe('user') expect(Array.isArray(lastMessage.content)).toBe(true) - const textContent = lastMessage.content.find((c: any) => c.type === 'text') as { type: 'text'; text: string } | undefined + const textContent = lastMessage.content.find( + (c: any) => c.type === 'text', + ) as { type: 'text'; text: string } | undefined expect(textContent).toBeDefined() // The text should be wrapped in tags @@ -754,12 +831,15 @@ describe('Run Cancellation Handling', () => { spyOn(databaseModule, 'finishAgentRun').mockResolvedValue(undefined) spyOn(databaseModule, 'addAgentStep').mockResolvedValue('step-1') - const serverSessionState = getInitialSessionState(getStubProjectFileContext()) + const serverSessionState = getInitialSessionState( + getStubProjectFileContext(), + ) serverSessionState.mainAgentState.messageHistory.push( userMessage('User prompt'), assistantMessage('Done!'), ) - const originalHistoryLength = serverSessionState.mainAgentState.messageHistory.length + const originalHistoryLength = + serverSessionState.mainAgentState.messageHistory.length spyOn(mainPromptModule, 'callMainPrompt').mockImplementation( async (params: Parameters[0]) => { @@ -823,7 +903,9 @@ describe('Run Cancellation Handling', () => { const abortController = new AbortController() // First run: server processes the user message and does some work, then user cancels - const firstRunServerState = getInitialSessionState(getStubProjectFileContext()) + const firstRunServerState = getInitialSessionState( + getStubProjectFileContext(), + ) firstRunServerState.mainAgentState.messageHistory.push( userMessage('Fix the bug in auth.ts'), assistantMessage('I will analyze the authentication module.'), @@ -847,7 +929,11 @@ describe('Run Cancellation Handling', () => { // Agent runtime adds interruption message on abort firstRunServerState.mainAgentState.messageHistory.push( - userMessage(withSystemTags("User interrupted the response. The assistant's previous work has been preserved.")) + userMessage( + withSystemTags( + "User interrupted the response. The assistant's previous work has been preserved.", + ), + ), ) // Server still sends the prompt-response with its session state @@ -886,12 +972,16 @@ describe('Run Cancellation Handling', () => { // Verify the first run preserved the user message and work expect(firstRunResult.sessionState).toBeDefined() - const firstHistory = firstRunResult.sessionState!.mainAgentState.messageHistory + const firstHistory = + firstRunResult.sessionState!.mainAgentState.messageHistory expect(firstHistory.length).toBe(3) // user + assistant + interruption const firstUserMsg = firstHistory.find( - (m) => m.role === 'user' && - m.content.some((c: any) => c.type === 'text' && c.text.includes('Fix the bug')) + (m) => + m.role === 'user' && + m.content.some( + (c: any) => c.type === 'text' && c.text.includes('Fix the bug'), + ), ) expect(firstUserMsg).toBeDefined() @@ -911,7 +1001,9 @@ describe('Run Cancellation Handling', () => { spyOn(databaseModule, 'addAgentStep').mockResolvedValue('step-2') // Second run: server receives the previous state and adds the new user message - const secondRunServerState = JSON.parse(JSON.stringify(firstRunResult.sessionState!)) as typeof firstRunServerState + const secondRunServerState = JSON.parse( + JSON.stringify(firstRunResult.sessionState!), + ) as typeof firstRunServerState secondRunServerState.mainAgentState.messageHistory.push( userMessage('Now also fix the login page'), assistantMessage('I will fix both issues.'), @@ -952,29 +1044,41 @@ describe('Run Cancellation Handling', () => { // Verify the second run's session state includes history from BOTH runs expect(secondRunResult.sessionState).toBeDefined() - const secondHistory = secondRunResult.sessionState!.mainAgentState.messageHistory + const secondHistory = + secondRunResult.sessionState!.mainAgentState.messageHistory // Should have: first user msg + first assistant msg + interruption + second user msg + second assistant msg expect(secondHistory.length).toBe(5) // The first user message should be present const firstUserMsgInSecond = secondHistory.find( - (m) => m.role === 'user' && - m.content.some((c: any) => c.type === 'text' && c.text.includes('Fix the bug')) + (m) => + m.role === 'user' && + m.content.some( + (c: any) => c.type === 'text' && c.text.includes('Fix the bug'), + ), ) expect(firstUserMsgInSecond).toBeDefined() // The second user message should also be present const secondUserMsg = secondHistory.find( - (m) => m.role === 'user' && - m.content.some((c: any) => c.type === 'text' && c.text.includes('fix the login page')) + (m) => + m.role === 'user' && + m.content.some( + (c: any) => + c.type === 'text' && c.text.includes('fix the login page'), + ), ) expect(secondUserMsg).toBeDefined() // The first assistant message should be preserved const firstAssistantMsg = secondHistory.find( - (m) => m.role === 'assistant' && - m.content.some((c: any) => c.type === 'text' && c.text.includes('authentication module')) + (m) => + m.role === 'assistant' && + m.content.some( + (c: any) => + c.type === 'text' && c.text.includes('authentication module'), + ), ) expect(firstAssistantMsg).toBeDefined() }) @@ -994,7 +1098,9 @@ describe('Run Cancellation Handling', () => { spyOn(databaseModule, 'addAgentStep').mockResolvedValue('step-1') const abortController = new AbortController() - const serverSessionState = getInitialSessionState(getStubProjectFileContext()) + const serverSessionState = getInitialSessionState( + getStubProjectFileContext(), + ) // Simulate multiple tool calls and results (more complex work done) serverSessionState.mainAgentState.messageHistory.push( @@ -1015,7 +1121,12 @@ describe('Run Cancellation Handling', () => { role: 'tool', toolCallId: 'read-1', toolName: 'read_files', - content: [{ type: 'json', value: [{ path: 'src/bug.ts', content: 'buggy code' }] }], + content: [ + { + type: 'json', + value: [{ path: 'src/bug.ts', content: 'buggy code' }], + }, + ], }, { role: 'assistant', @@ -1033,7 +1144,12 @@ describe('Run Cancellation Handling', () => { role: 'tool', toolCallId: 'write-1', toolName: 'write_file', - content: [{ type: 'json', value: { file: 'src/bug.ts', message: 'File written' } }], + content: [ + { + type: 'json', + value: { file: 'src/bug.ts', message: 'File written' }, + }, + ], }, ) @@ -1059,7 +1175,11 @@ describe('Run Cancellation Handling', () => { // Simulate agent runtime adding interruption message on abort serverSessionState.mainAgentState.messageHistory.push( - userMessage(withSystemTags("User interrupted the response. The assistant's previous work has been preserved.")) + userMessage( + withSystemTags( + "User interrupted the response. The assistant's previous work has been preserved.", + ), + ), ) // Server still returns the full session state @@ -1117,6 +1237,8 @@ describe('Run Cancellation Handling', () => { // Verify interruption message was added at the end const lastMessage = messageHistory[messageHistory.length - 1] expect(lastMessage.role).toBe('user') - expect((lastMessage.content[0] as { type: 'text'; text: string }).text).toContain('User interrupted the response') + expect( + (lastMessage.content[0] as { type: 'text'; text: string }).text, + ).toContain('User interrupted the response') }) }) diff --git a/sdk/src/run.ts b/sdk/src/run.ts index 2dfcef553..8d0c7986f 100644 --- a/sdk/src/run.ts +++ b/sdk/src/run.ts @@ -538,7 +538,13 @@ async function runOnce({ error && typeof error === 'object' && 'responseBody' in error ? (error as { responseBody: unknown }).responseBody : undefined - const { errorCode, message: parsedMessage } = parseApiErrorResponseBody(responseBody) + const { + countryBlockReason, + countryCode, + errorCode, + ipPrivacySignals, + message: parsedMessage, + } = parseApiErrorResponseBody(responseBody) if (parsedMessage) { errorMessage = parsedMessage } @@ -550,6 +556,9 @@ async function runOnce({ message: errorMessage, ...(statusCode !== undefined && { statusCode }), ...(errorCode !== undefined && { error: errorCode }), + ...(countryCode !== undefined && { countryCode }), + ...(countryBlockReason !== undefined && { countryBlockReason }), + ...(ipPrivacySignals !== undefined && { ipPrivacySignals }), }, }) }) diff --git a/web/src/app/api/v1/chat/completions/__tests__/completions.test.ts b/web/src/app/api/v1/chat/completions/__tests__/completions.test.ts index f12362ab6..ee66f6919 100644 --- a/web/src/app/api/v1/chat/completions/__tests__/completions.test.ts +++ b/web/src/app/api/v1/chat/completions/__tests__/completions.test.ts @@ -593,6 +593,7 @@ describe('/api/v1/chat/completions POST endpoint', () => { const body = await response.json() expect(body.error).toBe('free_mode_unavailable') expect(body.countryCode).toBe('UNKNOWN') + expect(body.countryBlockReason).toBe('missing_client_ip') }) it('rejects free-mode requests from anonymized Cloudflare country codes', async () => { @@ -634,6 +635,7 @@ describe('/api/v1/chat/completions POST endpoint', () => { const body = await response.json() expect(body.error).toBe('free_mode_unavailable') expect(body.countryCode).toBe('UNKNOWN') + expect(body.countryBlockReason).toBe('anonymized_or_unknown_country') }) it('lets freebuff use GLM 5.1 through Fireworks availability rules', async () => { diff --git a/web/src/app/api/v1/chat/completions/_post.ts b/web/src/app/api/v1/chat/completions/_post.ts index 84943dbf6..ca252682f 100644 --- a/web/src/app/api/v1/chat/completions/_post.ts +++ b/web/src/app/api/v1/chat/completions/_post.ts @@ -292,6 +292,8 @@ export async function postChatCompletions(params: { error: 'free_mode_unavailable', message: 'Free mode is not available in your country.', countryCode: countryAccess.countryCode ?? 'UNKNOWN', + countryBlockReason: countryAccess.blockReason, + ipPrivacySignals: countryAccess.ipPrivacy?.signals, }, { status: 403 }, ) diff --git a/web/src/app/api/v1/freebuff/session/__tests__/session.test.ts b/web/src/app/api/v1/freebuff/session/__tests__/session.test.ts index a7eaaa7cd..3e08ef944 100644 --- a/web/src/app/api/v1/freebuff/session/__tests__/session.test.ts +++ b/web/src/app/api/v1/freebuff/session/__tests__/session.test.ts @@ -165,6 +165,7 @@ describe('POST /api/v1/freebuff/session', () => { const body = await resp.json() expect(body.status).toBe('country_blocked') expect(body.countryCode).toBe('FR') + expect(body.countryBlockReason).toBe('country_not_allowed') expect(sessionDeps.rows.size).toBe(0) }) @@ -178,6 +179,7 @@ describe('POST /api/v1/freebuff/session', () => { const body = await resp.json() expect(body.status).toBe('country_blocked') expect(body.countryCode).toBe('UNKNOWN') + expect(body.countryBlockReason).toBe('missing_client_ip') expect(sessionDeps.rows.size).toBe(0) }) @@ -191,6 +193,7 @@ describe('POST /api/v1/freebuff/session', () => { const body = await resp.json() expect(body.status).toBe('country_blocked') expect(body.countryCode).toBe('UNKNOWN') + expect(body.countryBlockReason).toBe('anonymized_or_unknown_country') expect(sessionDeps.rows.size).toBe(0) }) @@ -256,6 +259,7 @@ describe('GET /api/v1/freebuff/session', () => { const body = await resp.json() expect(body.status).toBe('country_blocked') expect(body.countryCode).toBe('FR') + expect(body.countryBlockReason).toBe('country_not_allowed') }) test('returns banned 403 on GET for banned user', async () => { diff --git a/web/src/app/api/v1/freebuff/session/_handlers.ts b/web/src/app/api/v1/freebuff/session/_handlers.ts index 716a8a3c2..3418f188b 100644 --- a/web/src/app/api/v1/freebuff/session/_handlers.ts +++ b/web/src/app/api/v1/freebuff/session/_handlers.ts @@ -34,6 +34,8 @@ async function countryBlockedResponse( { status: 'country_blocked', countryCode: countryAccess.countryCode ?? 'UNKNOWN', + countryBlockReason: countryAccess.blockReason, + ipPrivacySignals: countryAccess.ipPrivacy?.signals, }, { status: 403 }, ) diff --git a/web/src/server/__tests__/free-mode-country.test.ts b/web/src/server/__tests__/free-mode-country.test.ts index ad3e57a5a..6026c3e01 100644 --- a/web/src/server/__tests__/free-mode-country.test.ts +++ b/web/src/server/__tests__/free-mode-country.test.ts @@ -124,6 +124,24 @@ describe('free mode country access', () => { expect(access.ipPrivacy?.signals).toEqual(['res_proxy']) }) + test('allows allowlisted countries when IPinfo only reports hosting or service', async () => { + const access = await getFreeModeCountryAccess( + makeReq({ + 'cf-ipcountry': 'US', + 'x-forwarded-for': '203.0.113.10', + }), + { + ipinfoToken: 'test-token', + lookupIpPrivacy: async () => ({ + signals: ['hosting', 'service'], + }), + }, + ) + expect(access.allowed).toBe(true) + expect(access.blockReason).toBe(null) + expect(access.ipPrivacy?.signals).toEqual(['hosting', 'service']) + }) + test('allows allowlisted countries when privacy lookup finds no anonymous signals', async () => { const access = await getFreeModeCountryAccess( makeReq({ @@ -204,4 +222,22 @@ describe('free mode country access', () => { signals: ['anonymous'], }) }) + + test('treats is_anonymous as blocking even when service is present', async () => { + const fetch = async () => + Response.json({ + service: 'Privacy Provider', + is_anonymous: true, + }) + + const privacy = await lookupIpinfoPrivacy({ + ip: '198.51.100.44', + token: 'test-token', + fetch: fetch as unknown as typeof globalThis.fetch, + }) + + expect(privacy).toEqual({ + signals: ['service', 'anonymous'], + }) + }) }) diff --git a/web/src/server/free-mode-country.ts b/web/src/server/free-mode-country.ts index 55490a6e1..84c210348 100644 --- a/web/src/server/free-mode-country.ts +++ b/web/src/server/free-mode-country.ts @@ -1,6 +1,10 @@ import geoip from 'geoip-lite' import type { NextRequest } from 'next/server' +import type { + FreebuffCountryBlockReason, + FreebuffIpPrivacySignal, +} from '@codebuff/common/types/freebuff-session' export const FREE_MODE_ALLOWED_COUNTRIES = new Set([ 'US', @@ -23,22 +27,8 @@ export const FREE_MODE_ALLOWED_COUNTRIES = new Set([ const CLOUDFLARE_ANONYMIZED_OR_UNKNOWN_COUNTRIES = new Set(['T1', 'XX']) -export type FreeModeCountryBlockReason = - | 'country_not_allowed' - | 'anonymized_or_unknown_country' - | 'anonymous_network' - | 'missing_client_ip' - | 'unresolved_client_ip' - -export type FreeModeIpPrivacySignal = - | 'anonymous' - | 'vpn' - | 'proxy' - | 'tor' - | 'relay' - | 'res_proxy' - | 'hosting' - | 'service' +export type FreeModeCountryBlockReason = FreebuffCountryBlockReason +export type FreeModeIpPrivacySignal = FreebuffIpPrivacySignal export type FreeModeIpPrivacy = { signals: FreeModeIpPrivacySignal[] @@ -78,6 +68,15 @@ const ipinfoPrivacyCache = new Map< { expiresAt: number; privacy: FreeModeIpPrivacy | null } >() +const FREE_MODE_BLOCKED_PRIVACY_SIGNALS = new Set([ + 'anonymous', + 'vpn', + 'proxy', + 'tor', + 'relay', + 'res_proxy', +]) + export function extractClientIp(req: NextRequest): string | undefined { const forwardedFor = req.headers.get('x-forwarded-for') if (forwardedFor) { @@ -135,7 +134,10 @@ function privacySignalsFromIpinfo( ) { signals.push('service') } - if (signals.length === 0 && data.is_anonymous === true) { + if ( + data.is_anonymous === true && + !signals.some((signal) => FREE_MODE_BLOCKED_PRIVACY_SIGNALS.has(signal)) + ) { signals.push('anonymous') } return signals @@ -268,7 +270,11 @@ export async function getFreeModeCountryAccess( } const ipPrivacy = await getIpPrivacy(clientIp, options) - if (ipPrivacy?.signals.length) { + if ( + ipPrivacy?.signals.some((signal) => + FREE_MODE_BLOCKED_PRIVACY_SIGNALS.has(signal), + ) + ) { return { ...baseAccess, allowed: false,