From 81eeb9add536d24c06274099af043e9a55a22285 Mon Sep 17 00:00:00 2001 From: tmm Date: Sun, 26 Apr 2026 16:58:30 -0400 Subject: [PATCH 1/2] feat: capture request ai-agent --- .gitattributes | 1 + cli/src/client.test.ts | 165 +++++++++++++++- cli/src/client.ts | 117 ++++++++--- db/codegen.ts | 185 ++++++++++++------ .../20260426000000_request_ai_agent.ts | 15 ++ db/schemas.gen.ts | 147 ++++++++++++++ db/types.gen.ts | 1 + plugins/amp/src/plugin.test.ts | 22 ++- plugins/amp/src/plugin.ts | 4 + plugins/claude/src/server.test.ts | 16 +- plugins/claude/src/server.ts | 4 + plugins/opencode/src/server.test.ts | 10 +- plugins/opencode/src/server.ts | 4 + plugins/pi/src/extension.test.ts | 40 ++-- plugins/pi/src/index.ts | 7 +- src/api.ts | 25 ++- src/og.tsx | 19 +- src/queues/request.ts | 2 + src/queues/request.workers.test.ts | 2 + 19 files changed, 647 insertions(+), 139 deletions(-) create mode 100644 db/migrations/20260426000000_request_ai_agent.ts create mode 100644 db/schemas.gen.ts diff --git a/.gitattributes b/.gitattributes index 9665d08e..bd2111d7 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1,5 +1,6 @@ pnpm-lock.yaml linguist-generated cli/src/cf-env.d.ts linguist-generated +db/schemas.gen.ts linguist-generated db/types.gen.ts linguist-generated src/routeTree.gen.ts linguist-generated src/worker-configuration.d.ts linguist-generated diff --git a/cli/src/client.test.ts b/cli/src/client.test.ts index bda752ce..c7bd895d 100644 --- a/cli/src/client.test.ts +++ b/cli/src/client.test.ts @@ -1,8 +1,12 @@ import { HttpResponse, http } from 'msw' -import { expect, test } from 'vitest' +import { afterEach, expect, test, vi } from 'vitest' import { server } from '../test/server.ts' import { createClient, defaultBaseUrl } from './client.ts' +afterEach(() => { + vi.unstubAllEnvs() +}) + test('createClient.fetch moves target fragment into anchor query', async () => { let requestUrl: URL | undefined server.use( @@ -56,3 +60,162 @@ test('createClient.fetch leaves hash-free target urls unchanged', async () => { expect(requestUrl?.pathname).toBe('/api/example.com/foo%3Ftab%3Dapi') expect(requestUrl?.searchParams.has('anchor')).toBe(false) }) + +test('createClient adds explicit ai agent header', async () => { + let aiAgent: string | null | undefined + server.use( + http.get('*', ({ request }) => { + aiAgent = request.headers.get('x-ai-agent') + return HttpResponse.json({ version: 'x.y.z', published_at: null }) + }), + ) + + const client = createClient(defaultBaseUrl, { aiAgent: 'gemini' }) + const res = await client.api.cli.latest.$get({ query: {} }) + + expect(res.status).toBe(200) + await expect(res.json()).resolves.toEqual({ published_at: null, version: 'x.y.z' }) + expect(aiAgent).toBe('gemini') +}) + +test('createClient aiAgent overrides std-env fallback detection', async () => { + let aiAgent: string | null | undefined + vi.stubEnv('AI_AGENT', 'codex') + server.use( + http.get('*', ({ request }) => { + aiAgent = request.headers.get('x-ai-agent') + return HttpResponse.json({ version: 'x.y.z', published_at: null }) + }), + ) + + const client = createClient(defaultBaseUrl, { aiAgent: 'gemini' }) + const res = await client.api.cli.latest.$get({ query: {} }) + + expect(res.status).toBe(200) + await expect(res.json()).resolves.toEqual({ published_at: null, version: 'x.y.z' }) + expect(aiAgent).toBe('gemini') +}) + +test('createClient omits unsupported explicit aiAgent values without std-env fallback', async () => { + const aiAgents: Array = [] + vi.stubEnv('AI_AGENT', 'codex') + server.use( + http.get('*', ({ request }) => { + aiAgents.push(request.headers.get('x-ai-agent')) + return HttpResponse.json({ version: 'x.y.z', published_at: null }) + }), + ) + + for (const aiAgent of ['', 'invalid']) { + const client = createClient(defaultBaseUrl, { aiAgent: aiAgent as never }) + const res = await client.api.cli.latest.$get({ query: {} }) + + expect(res.status).toBe(200) + await expect(res.json()).resolves.toEqual({ published_at: null, version: 'x.y.z' }) + } + + expect(aiAgents).toEqual([null, null]) +}) + +test('createClient falls back to std-env ai agent detection', async () => { + let aiAgent: string | null | undefined + vi.stubEnv('AI_AGENT', 'codex') + server.use( + http.get('*', ({ request }) => { + aiAgent = request.headers.get('x-ai-agent') + return HttpResponse.json({ version: 'x.y.z', published_at: null }) + }), + ) + + const client = createClient(defaultBaseUrl) + const res = await client.api.cli.latest.$get({ query: {} }) + + expect(res.status).toBe(200) + await expect(res.json()).resolves.toEqual({ published_at: null, version: 'x.y.z' }) + expect(aiAgent).toBe('codex') +}) + +test('createClient omits ai agent header when std-env detects an unsupported agent', async () => { + let aiAgent: string | null | undefined + vi.stubEnv('AI_AGENT', 'devin') + server.use( + http.get('*', ({ request }) => { + aiAgent = request.headers.get('x-ai-agent') + return HttpResponse.json({ version: 'x.y.z', published_at: null }) + }), + ) + + const client = createClient(defaultBaseUrl) + const res = await client.api.cli.latest.$get({ query: {} }) + + expect(res.status).toBe(200) + await expect(res.json()).resolves.toEqual({ published_at: null, version: 'x.y.z' }) + expect(aiAgent).toBeNull() +}) + +test('createClient.fetch keeps ai agent header when adding authorization', async () => { + let aiAgent: string | null | undefined + let authorization: string | null | undefined + server.use( + http.get('*', ({ request }) => { + aiAgent = request.headers.get('x-ai-agent') + authorization = request.headers.get('authorization') + return HttpResponse.json({ content: '# Example' }) + }), + ) + + const client = createClient(defaultBaseUrl, { aiAgent: 'gemini' }) + const res = await client.fetch('example.com', { token: 'curlmd_test' }) + + expect(res.status).toBe(200) + await expect(res.json()).resolves.toEqual({ content: '# Example' }) + expect(aiAgent).toBe('gemini') + expect(authorization).toBe('Bearer curlmd_test') +}) + +test('createClient omits ai agent header in browser-like environment without process', async () => { + let aiAgent: string | null | undefined + const originalProcess = globalThis.process + // @ts-expect-error -- simulate browser environment + globalThis.process = undefined + server.use( + http.get('*', ({ request }) => { + aiAgent = request.headers.get('x-ai-agent') + return HttpResponse.json({ version: 'x.y.z', published_at: null }) + }), + ) + + try { + const client = createClient(defaultBaseUrl) + const res = await client.api.cli.latest.$get({ query: {} }) + + expect(res.status).toBe(200) + await expect(res.json()).resolves.toEqual({ published_at: null, version: 'x.y.z' }) + expect(aiAgent).toBeNull() + } finally { + globalThis.process = originalProcess + } +}) + +test('createClient keeps ai agent header when route requests add headers', async () => { + let aiAgent: string | null | undefined + let organizationId: string | null | undefined + server.use( + http.get('*', ({ request }) => { + aiAgent = request.headers.get('x-ai-agent') + organizationId = request.headers.get('x-organization-id') + return HttpResponse.json({ version: 'x.y.z', published_at: null }) + }), + ) + + const client = createClient(defaultBaseUrl, { aiAgent: 'gemini' }) + const res = await client.api.cli.latest.$get( + { query: {} }, + { headers: { 'x-organization-id': 'org_123' } }, + ) + + expect(res.status).toBe(200) + await expect(res.json()).resolves.toEqual({ published_at: null, version: 'x.y.z' }) + expect(aiAgent).toBe('gemini') + expect(organizationId).toBe('org_123') +}) diff --git a/cli/src/client.ts b/cli/src/client.ts index 60cdd758..bd51e7bf 100644 --- a/cli/src/client.ts +++ b/cli/src/client.ts @@ -17,15 +17,29 @@ export const defaultBaseUrl = 'https://curl.md' * const res = await client.fetch('example.com') * ``` */ -export function createClient(url: string = defaultBaseUrl, options?: ClientRequestOptions): Client { - const client = hc(url, options) +export function createClient( + url: string = defaultBaseUrl, + options?: ClientRequestOptions & { + aiAgent?: AiAgent | undefined + }, +): Client { + const aiAgent = (() => { + if (options && Object.prototype.hasOwnProperty.call(options, 'aiAgent')) + return normalizeAiAgent(options.aiAgent) + return normalizeAiAgent(detectAgent().name) + })() + + const client = hc(url, { + ...options, + headers: getHeaders(options?.headers, undefined, { aiAgent }), + }) return new Proxy(client, { get(target, prop, receiver) { if (prop === 'fetch') { return (targetUrl: string, fetchOptions?: FetchOptions | undefined) => { const normalizedTargetURL = normalizeTargetURL(targetUrl) - const { options, token, ...queryOptions } = fetchOptions ?? {} + const { options: requestOptions, token, ...queryOptions } = fetchOptions ?? {} const query = { anchor: normalizedTargetURL.anchor, ...queryOptions, @@ -35,12 +49,20 @@ export function createClient(url: string = defaultBaseUrl, options?: ClientReque anchor?: string | undefined } + const clientRequestOptions = (() => { + if (!requestOptions && !token) return undefined + return { + ...requestOptions, + headers: getHeaders(options?.headers, requestOptions?.headers, { aiAgent, token }), + } + })() + return target.api[':url{.+}'].$get( { param: { url: normalizedTargetURL.url }, query, }, - token ? withAuthorizationHeader(options, token) : options, + clientRequestOptions, ) } } @@ -69,28 +91,77 @@ type FetchOptions = Partial> & { token?: string | undefined } -function withAuthorizationHeader( - options: NonNullable[1]> | undefined, - token: string, -): NonNullable[1]> { - const headers = options?.headers - return { - ...options, - headers: - typeof headers === 'function' - ? async () => withTokenHeader(await headers(), token) - : withTokenHeader(headers, token), - } +function getHeaders( + baseHeaders: ClientRequestOptions['headers'], + requestHeaders: ClientRequestOptions['headers'], + opts: { + aiAgent?: AiAgent | undefined + token?: string | undefined + }, +): ClientRequestOptions['headers'] { + if (typeof baseHeaders === 'function' || typeof requestHeaders === 'function') + return async () => { + const nextBaseHeaders = typeof baseHeaders === 'function' ? await baseHeaders() : baseHeaders + const nextRequestHeaders = + typeof requestHeaders === 'function' ? await requestHeaders() : requestHeaders + return mergeHeaders(nextBaseHeaders, nextRequestHeaders, opts) || {} + } + return mergeHeaders(baseHeaders, requestHeaders, opts) +} + +function mergeHeaders( + baseHeaders: Record | undefined, + requestHeaders: Record | undefined, + opts: { + aiAgent?: AiAgent | undefined + token?: string | undefined + }, +) { + if (!baseHeaders && !requestHeaders && !opts.aiAgent && !opts.token) return undefined + const nextHeaders = new Headers(baseHeaders) + if (requestHeaders) + for (const [name, value] of Object.entries(requestHeaders)) nextHeaders.set(name, value) + nextHeaders.delete('x-ai-agent') + if (opts.aiAgent) nextHeaders.set('x-ai-agent', opts.aiAgent) + if (opts.token) nextHeaders.set('Authorization', `Bearer ${opts.token}`) + return Object.fromEntries(nextHeaders.entries()) } -function withTokenHeader(headers: Record | undefined, token: string) { - const nextHeaders = { ...headers } - for (const key of Object.keys(nextHeaders)) { - if (key.toLowerCase() !== 'authorization') continue - delete nextHeaders[key] +const aiAgents = ['amp', 'claude', 'codex', 'cursor', 'gemini', 'opencode', 'pi'] as const +type AiAgent = (typeof aiAgents)[number] + +function normalizeAiAgent(value: unknown): AiAgent | undefined { + if (typeof value !== 'string') return undefined + + const normalizedValue = value.toLowerCase() + if (!aiAgents.includes(normalizedValue as AiAgent)) return undefined + return normalizedValue as AiAgent +} + +/** + * Vendored from std-env v4.1.0 — agent detection subset + * https://github.com/unjs/std-env + * MIT License + */ +function detectAgent(): { name?: string } { + const env = globalThis.process?.env || Object.create(null) + const aiAgent = env.AI_AGENT + if (aiAgent) return { name: aiAgent.toLowerCase() } + + const rules: Array<[string, (string | (() => boolean))[]]> = [ + ['amp', [() => env.AGENT === 'amp']], + ['claude', ['CLAUDECODE', 'CLAUDE_CODE']], + ['codex', ['CODEX_SANDBOX', 'CODEX_THREAD_ID']], + ['cursor', ['CURSOR_AGENT']], + ['gemini', ['GEMINI_CLI']], + ['opencode', ['OPENCODE']], + ['pi', [() => /\.pi[\\/]agent/.test(env.PATH ?? '')]], + ] + for (const [name, checks] of rules) { + for (const check of checks) + if (typeof check === 'string' ? env[check] : check()) return { name } } - nextHeaders.Authorization = `Bearer ${token}` - return nextHeaders + return {} } function normalizeTargetURL(targetUrl: string) { diff --git a/db/codegen.ts b/db/codegen.ts index 98bed6f0..3263e6ec 100644 --- a/db/codegen.ts +++ b/db/codegen.ts @@ -45,121 +45,176 @@ const customTypes: Record> = { organization_invite: { role: "'admin' | 'member' | 'owner'" }, organization_member: { role: "'admin' | 'member' | 'owner'" }, request: { + ai_agent: "'amp' | 'claude' | 'codex' | 'cursor' | 'gemini' | 'opencode' | 'pi'", mode: "'rush' | 'smart'", source_tokens_method: "'estimated' | 'html' | 'markdown'", }, session: { session_type: "'browser' | 'cli'" }, } -const timestampTypes = new Set([ - 'timestamptz', - 'timestamp', - 'timestamp with time zone', - 'timestamp without time zone', -]) - -function pgToTs(dataType: string): string { - switch (dataType) { - case 'varchar': - case 'text': - case 'character varying': - return 'string' - case 'int4': - case 'integer': - case 'int': - case 'smallint': - case 'bigint': - case 'float4': - case 'float8': - case 'real': - case 'double precision': - case 'numeric': - return 'number' - case 'bool': - case 'boolean': - return 'boolean' - case 'bytea': - return 'Uint8Array' - default: - return 'unknown' - } -} - let output = '// Auto-generated from database schema\n\n' output += "import type * as k from 'kysely'\n\n" output += 'type Timestamp = k.ColumnType\n' output += 'type GeneratedTimestamp = k.ColumnType\n\n' +let schemaOutput = '// Auto-generated from database schema\n\n' +schemaOutput += "import { z } from 'zod'\n\n" + +// TODO: Generate insert/update schemas alongside row/select schemas. + output += 'export interface DB {\n' -for (const table of publicTables) { - output += `\t${table.name}: ${table.name}\n` -} +for (const table of publicTables) output += `\t${table.name}: ${table.name}\n` output += '}\n\n' +const timestampTypes = new Set([ + 'timestamptz', + 'timestamp', + 'timestamp with time zone', + 'timestamp without time zone', +]) + for (const table of publicTables) { const columns = [...table.columns].sort((a, b) => a.name.localeCompare(b.name)) + const tableSchemaName = table.name + let tableSchemaObjectOutput = `export const ${tableSchemaName} = z.object({\n` output += `type ${table.name} = {\n` for (const col of columns) { const custom = customTypes[table.name]?.[col.name] const isTimestamp = timestampTypes.has(col.dataType) const enumValues = enums.get(`public.${col.dataType}`) + const customValues = custom ? parseStringUnionValues(custom) : undefined // Timestamp columns use rich ColumnType aliases if (isTimestamp) { const base = col.hasDefaultValue ? 'GeneratedTimestamp' : 'Timestamp' const suffix = col.isNullable ? ' | null' : '' output += `\t${col.name}: ${base}${suffix}\n` + tableSchemaObjectOutput += ` ${col.name}: z.date()${col.isNullable ? '.nullable()' : ''},\n` continue } // Custom override, auto-discovered enum, or standard type mapping - let baseType: string - if (custom) baseType = custom - else if (enumValues) - baseType = enumValues - .sort() - .map((v) => `'${v}'`) - .join(' | ') - else baseType = pgToTs(col.dataType) + const baseType = (() => { + if (custom) return custom + else if (enumValues) + return enumValues + .sort() + .map((v) => `'${v}'`) + .join(' | ') + return pgToTs(col.dataType) + })() + + const zodType = (() => { + if (customValues) return `z.enum([${customValues.map((value) => `'${value}'`).join(', ')}])` + if (enumValues) { + const sortedValues = [...enumValues].sort() + return `z.enum([${sortedValues.map((value) => `'${value}'`).join(', ')}])` + } + return pgToZod(col.dataType) + })() const nullableSuffix = col.isNullable ? ' | null' : '' const isGenerated = col.hasDefaultValue - if (isGenerated) { - output += `\t${col.name}: k.Generated<${baseType}${nullableSuffix}>\n` - } else { - output += `\t${col.name}: ${baseType}${nullableSuffix}\n` - } + if (isGenerated) output += `\t${col.name}: k.Generated<${baseType}${nullableSuffix}>\n` + else output += `\t${col.name}: ${baseType}${nullableSuffix}\n` + + tableSchemaObjectOutput += ` ${col.name}: ${zodType}${col.isNullable ? '.nullable()' : ''},\n` } output += '}\n\n' + tableSchemaObjectOutput += '})\n' + schemaOutput += `${tableSchemaObjectOutput}\n` } output += 'export declare namespace DB {\n' -for (const table of publicTables) { +for (const table of publicTables) output += `\ttype ${table.name} = k.Selectable\n` -} output += '\n\texport namespace Insertable {\n' -for (const table of publicTables) { +for (const table of publicTables) output += `\t\ttype ${table.name} = k.Insertable\n` -} output += '\t}\n\n\texport namespace Selectable {\n' -for (const table of publicTables) { +for (const table of publicTables) output += `\t\ttype ${table.name} = k.Selectable\n` -} output += '\t}\n\n\texport namespace Updateable {\n' -for (const table of publicTables) { +for (const table of publicTables) output += `\t\ttype ${table.name} = k.Updateable\n` -} output += '\t}\n}\n' -const outputPath = path.resolve(import.meta.dirname, '../db/types.gen.ts') -fs.writeFileSync(outputPath, `${output.trimEnd()}\n`) -execSync(`pnpm exec oxfmt ${outputPath}`, { - cwd: path.resolve(import.meta.dirname, '..'), - stdio: 'inherit', -}) -console.log('Generated db/types.gen.ts') +schemaOutput += 'export const db = {\n' +for (const table of publicTables) schemaOutput += ` ${table.name}: ${table.name},\n` +schemaOutput += '}\n' + +writeGeneratedFile('db/types.gen.ts', output) +writeGeneratedFile('db/schemas.gen.ts', schemaOutput) process.exit() + +function pgToTs(dataType: string): string { + switch (dataType) { + case 'varchar': + case 'text': + case 'character varying': + return 'string' + case 'int4': + case 'integer': + case 'int': + case 'smallint': + case 'bigint': + case 'float4': + case 'float8': + case 'real': + case 'double precision': + case 'numeric': + return 'number' + case 'bool': + case 'boolean': + return 'boolean' + case 'bytea': + return 'Uint8Array' + default: + return 'unknown' + } +} + +function pgToZod(dataType: string): string { + switch (dataType) { + case 'varchar': + case 'text': + case 'character varying': + return 'z.string()' + case 'int4': + case 'integer': + case 'int': + case 'smallint': + case 'bigint': + case 'float4': + case 'float8': + case 'real': + case 'double precision': + case 'numeric': + return 'z.number()' + case 'bool': + case 'boolean': + return 'z.boolean()' + case 'bytea': + return 'z.instanceof(Uint8Array)' + default: + return 'z.unknown()' + } +} + +function parseStringUnionValues(value: string) { + return [...value.matchAll(/'([^']+)'/g)].map((match) => match[1]!) +} + +function writeGeneratedFile(filePath: string, output: string) { + const outputPath = path.resolve(import.meta.dirname, '..', filePath) + fs.writeFileSync(outputPath, `${output.trimEnd()}\n`) + execSync(`pnpm exec oxfmt ${outputPath}`, { + cwd: path.resolve(import.meta.dirname, '..'), + stdio: 'inherit', + }) + console.log(`Generated ${filePath}`) +} diff --git a/db/migrations/20260426000000_request_ai_agent.ts b/db/migrations/20260426000000_request_ai_agent.ts new file mode 100644 index 00000000..a2167c88 --- /dev/null +++ b/db/migrations/20260426000000_request_ai_agent.ts @@ -0,0 +1,15 @@ +import { type Kysely, sql } from 'kysely' + +export async function up(db: Kysely): Promise { + await db.schema.alterTable('request').addColumn('ai_agent', 'text').execute() + + await sql`ALTER TABLE request ADD CONSTRAINT request_ai_agent_chk CHECK (ai_agent IN ('amp', 'claude', 'codex', 'cursor', 'gemini', 'opencode', 'pi'))`.execute( + db, + ) +} + +export async function down(db: Kysely): Promise { + await sql`ALTER TABLE request DROP CONSTRAINT IF EXISTS request_ai_agent_chk`.execute(db) + + await db.schema.alterTable('request').dropColumn('ai_agent').execute() +} diff --git a/db/schemas.gen.ts b/db/schemas.gen.ts new file mode 100644 index 00000000..1085c155 --- /dev/null +++ b/db/schemas.gen.ts @@ -0,0 +1,147 @@ +// Auto-generated from database schema + +import { z } from 'zod' + +export const account = z.object({ + avatar_url: z.string().nullable(), + balance_mills: z.number(), + created_at: z.date(), + default_payment_method_id: z.string().nullable(), + deleted_at: z.date().nullable(), + email: z.string(), + id: z.string(), + login: z.string(), + name: z.string().nullable(), + role: z.enum(['crew', 'user']), + stripe_customer_id: z.string().nullable(), +}) + +export const account_provider = z.object({ + access_token: z.string().nullable(), + access_token_expires_at: z.date().nullable(), + account_id: z.string(), + created_at: z.date(), + id: z.string(), + provider: z.string(), + provider_account_id: z.string(), + refresh_token: z.string().nullable(), + refresh_token_expires_at: z.date().nullable(), +}) + +export const api_key = z.object({ + account_id: z.string(), + created_at: z.date(), + deleted_at: z.date().nullable(), + id: z.string(), + key_hash: z.string(), + key_prefix: z.string(), + last_used_at: z.date().nullable(), + name: z.string(), + organization_id: z.string().nullable(), +}) + +export const credit_transaction = z.object({ + account_id: z.string().nullable(), + amount_mills: z.number(), + balance_after_mills: z.number(), + created_at: z.date(), + id: z.string(), + organization_id: z.string().nullable(), + reference_id: z.string().nullable(), + type: z.enum(['chargeback', 'promo', 'purchase', 'refund', 'request']), +}) + +export const device_code = z.object({ + account_id: z.string().nullable(), + code: z.string(), + created_at: z.date(), + expires_at: z.date(), + id: z.string(), + status: z.enum(['approved', 'pending']), + user_code: z.string(), +}) + +export const organization = z.object({ + balance_mills: z.number(), + created_at: z.date(), + default_payment_method_id: z.string().nullable(), + deleted_at: z.date().nullable(), + id: z.string(), + login: z.string(), + name: z.string(), + stripe_customer_id: z.string().nullable(), +}) + +export const organization_invite = z.object({ + created_at: z.date(), + created_by: z.string(), + deleted_at: z.date().nullable(), + expires_at: z.date(), + id: z.string(), + max_uses: z.number().nullable(), + organization_id: z.string(), + role: z.enum(['admin', 'member', 'owner']), + token: z.string(), + use_count: z.number(), +}) + +export const organization_member = z.object({ + account_id: z.string(), + created_at: z.date(), + id: z.string(), + organization_id: z.string(), + role: z.enum(['admin', 'member', 'owner']), +}) + +export const request = z.object({ + account_id: z.string().nullable(), + ai_agent: z.enum(['amp', 'claude', 'codex', 'cursor', 'gemini', 'opencode', 'pi']).nullable(), + api_key_id: z.string().nullable(), + cached: z.boolean(), + created_at: z.date(), + extracted_tokens: z.number().nullable(), + filtered_tokens: z.number().nullable(), + hostname: z.string(), + id: z.string(), + keywords: z.string().nullable(), + markdown_tokens: z.number(), + mode: z.enum(['rush', 'smart']).nullable(), + objective: z.string().nullable(), + organization_id: z.string().nullable(), + path: z.string(), + source_tokens: z.number(), + source_tokens_method: z.enum(['estimated', 'html', 'markdown']), + url: z.string(), + user_agent: z.string().nullable(), +}) + +export const session = z.object({ + account_id: z.string(), + created_at: z.date(), + expires_at: z.date(), + id: z.string(), + refresh_token_hash: z.string().nullable(), + session_type: z.enum(['browser', 'cli']), +}) + +export const session_access_token = z.object({ + created_at: z.date(), + expires_at: z.date(), + id: z.string(), + session_id: z.string(), + token_hash: z.string(), +}) + +export const db = { + account: account, + account_provider: account_provider, + api_key: api_key, + credit_transaction: credit_transaction, + device_code: device_code, + organization: organization, + organization_invite: organization_invite, + organization_member: organization_member, + request: request, + session: session, + session_access_token: session_access_token, +} diff --git a/db/types.gen.ts b/db/types.gen.ts index df7c8f97..68a80be8 100644 --- a/db/types.gen.ts +++ b/db/types.gen.ts @@ -112,6 +112,7 @@ type organization_member = { type request = { account_id: string | null + ai_agent: 'amp' | 'claude' | 'codex' | 'cursor' | 'gemini' | 'opencode' | 'pi' | null api_key_id: string | null cached: boolean created_at: GeneratedTimestamp diff --git a/plugins/amp/src/plugin.test.ts b/plugins/amp/src/plugin.test.ts index 0a156515..53ff7a29 100644 --- a/plugins/amp/src/plugin.test.ts +++ b/plugins/amp/src/plugin.test.ts @@ -278,10 +278,12 @@ test('prefers CURLMD_API_KEY for authentication', async () => { const { tools } = loadPlugin() await tools[0]!.execute({ url: 'https://example.com' }, { logger: { log() {} } } as any) - expect(requests[0]?.headers).toEqual({ - accept: 'application/json', - authorization: 'Bearer curlmd_test_token', - }) + expect(requests[0]?.headers).toEqual( + expect.objectContaining({ + accept: 'application/json', + authorization: 'Bearer curlmd_test_token', + }), + ) }) // --- Session auth --- @@ -316,11 +318,13 @@ test('uses session auth headers when available', async () => { const { tools } = loadPlugin() await tools[0]!.execute({ url: 'https://example.com' }, { logger: { log() {} } } as any) - expect(requests[0]?.headers).toEqual({ - accept: 'application/json', - authorization: 'Bearer access-token-1', - 'x-organization-id': 'org_123', - }) + expect(requests[0]?.headers).toEqual( + expect.objectContaining({ + accept: 'application/json', + authorization: 'Bearer access-token-1', + 'x-organization-id': 'org_123', + }), + ) Session.delete() }) diff --git a/plugins/amp/src/plugin.ts b/plugins/amp/src/plugin.ts index 02fc5ae0..af53a763 100644 --- a/plugins/amp/src/plugin.ts +++ b/plugins/amp/src/plugin.ts @@ -3,6 +3,8 @@ import type { PluginAPI } from '@ampcode/plugin' import { createClient, defaultBaseUrl } from 'curl.md' import { Auth, Session } from 'curl.md/internal' +const aiAgent = 'amp' as const + export default function (amp: PluginAPI) { const baseUrl = process.env.CURLMD_BASE_URL || defaultBaseUrl const apiKey = process.env.CURLMD_API_KEY @@ -110,6 +112,7 @@ export default function (amp: PluginAPI) { } const client = createClient(baseUrl, { + aiAgent, headers: apiKey ? createHeaders(null) : createHeaders(authHeaders), }) let res = await client.fetch(url, { ...fetchParams, token: apiKey }) @@ -118,6 +121,7 @@ export default function (amp: PluginAPI) { authHeaders = await resolver({ forceRefresh: true }) if (!authHeaders) authType = 'anon' const retryClient = createClient(baseUrl, { + aiAgent, headers: apiKey ? createHeaders(null) : createHeaders(authHeaders), }) res = await retryClient.fetch(url, { ...fetchParams, token: apiKey }) diff --git a/plugins/claude/src/server.test.ts b/plugins/claude/src/server.test.ts index 68b44ec8..37583ff5 100644 --- a/plugins/claude/src/server.test.ts +++ b/plugins/claude/src/server.test.ts @@ -117,7 +117,11 @@ test('returns stripped markdown and forwards fetch options', async () => { expect(requests[0]?.url).toContain('keywords=plugin') expect(requests[0]?.url).toContain('mode=smart') expect(requests[0]?.url).toContain('objective=Summarize+the+docs') - expect(requests[0]?.headers).toEqual({ accept: 'application/json' }) + expect(requests[0]?.headers).toEqual( + expect.objectContaining({ + accept: 'application/json', + }), + ) expect(result).toEqual({ content: [{ text: '# Example\n\n---\n\nPowered by [curl.md](https://curl.md)', type: 'text' }], }) @@ -249,10 +253,12 @@ test('uses CURLMD_API_KEY and surfaces invalid API key guidance on 401', async ( const tool = await loadTool() const result = await tool.handler({ url: 'https://example.com' }) - expect(requests[0]?.headers).toEqual({ - accept: 'application/json', - authorization: 'Bearer curlmd_test_token', - }) + expect(requests[0]?.headers).toEqual( + expect.objectContaining({ + accept: 'application/json', + authorization: 'Bearer curlmd_test_token', + }), + ) expectErrorText(result).toBe('curl.md authentication failed. Fix CURLMD_API_KEY.') }) diff --git a/plugins/claude/src/server.ts b/plugins/claude/src/server.ts index 8b811a8a..f463ba03 100644 --- a/plugins/claude/src/server.ts +++ b/plugins/claude/src/server.ts @@ -4,6 +4,8 @@ import { createClient, defaultBaseUrl } from 'curl.md' import { Auth, Session } from 'curl.md/internal' import { z } from 'zod' +const aiAgent = 'claude' as const + const baseUrl = process.env.CURLMD_BASE_URL || defaultBaseUrl const apiKey = process.env.CURLMD_API_KEY const resolver = Auth.createResolver(baseUrl, apiKey) @@ -73,6 +75,7 @@ async function fetchPage(input: { })() const client = createClient(baseUrl, { + aiAgent, headers: apiKey ? createHeaders(null) : createHeaders(authHeaders), }) let res = await client.fetch(url, { @@ -87,6 +90,7 @@ async function fetchPage(input: { authHeaders = await resolver({ forceRefresh: true }) if (!authHeaders) authType = 'anon' const retryClient = createClient(baseUrl, { + aiAgent, headers: apiKey ? createHeaders(null) : createHeaders(authHeaders), }) res = await retryClient.fetch(url, { diff --git a/plugins/opencode/src/server.test.ts b/plugins/opencode/src/server.test.ts index 99326d13..36fadbca 100644 --- a/plugins/opencode/src/server.test.ts +++ b/plugins/opencode/src/server.test.ts @@ -371,10 +371,12 @@ test('uses CURLMD_API_KEY when provided', async () => { const hooks = await loadPlugin({ webfetch: true }) await hooks.tool!.webfetch.execute({ url: 'https://example.com' }, createToolContext(vi.fn())) - expect(requests[0]?.headers).toEqual({ - accept: 'application/json', - authorization: 'Bearer curlmd_test_token', - }) + expect(requests[0]?.headers).toEqual( + expect.objectContaining({ + accept: 'application/json', + authorization: 'Bearer curlmd_test_token', + }), + ) }) test('surfaces authentication guidance for anonymous 401s', async () => { diff --git a/plugins/opencode/src/server.ts b/plugins/opencode/src/server.ts index e69c35ba..44f072df 100644 --- a/plugins/opencode/src/server.ts +++ b/plugins/opencode/src/server.ts @@ -3,6 +3,8 @@ import * as curlmd from 'curl.md' import * as curlmdInternal from 'curl.md/internal' import { createHeaders, formatApiError, parseApiError } from './utils.ts' +const aiAgent = 'opencode' as const + export const plugin: opencodePlugin.Plugin = async (_input, options) => { const baseUrl = process.env.CURLMD_BASE_URL || curlmd.defaultBaseUrl const apiKey = process.env.CURLMD_API_KEY @@ -154,6 +156,7 @@ async function fetchPage(input: { const apiKey = process.env.CURLMD_API_KEY const client = curlmd.createClient(input.baseUrl, { + aiAgent, headers: apiKey ? createHeaders(null) : createHeaders(authHeaders), }) let res = await client.fetch(url, { @@ -166,6 +169,7 @@ async function fetchPage(input: { authHeaders = await input.resolver({ forceRefresh: true }) if (!authHeaders) authType = 'anon' const retryClient = curlmd.createClient(input.baseUrl, { + aiAgent, headers: apiKey ? createHeaders(null) : createHeaders(authHeaders), }) res = await retryClient.fetch(url, { diff --git a/plugins/pi/src/extension.test.ts b/plugins/pi/src/extension.test.ts index a159e482..99753e1b 100644 --- a/plugins/pi/src/extension.test.ts +++ b/plugins/pi/src/extension.test.ts @@ -227,10 +227,12 @@ test('status command verifies API key auth state', async () => { expect(requests[0]?.url).toBe(`${defaultBaseUrl}/api/auth/me`) expect(requests[0]?.method).toBe('GET') - expect(requests[0]?.headers).toEqual({ - accept: 'application/json', - authorization: 'Bearer curlmd_test_token', - }) + expect(requests[0]?.headers).toEqual( + expect.objectContaining({ + accept: 'application/json', + authorization: 'Bearer curlmd_test_token', + }), + ) expect(notify).toHaveBeenCalledWith( `${extensionHeader}\nAuth: api_key (tmm)\nOrganization: none\nTool: read_web_page (alias: curl_md)\nCLI: ${mockCliPath}`, 'info', @@ -665,9 +667,11 @@ test('fetches markdown from curl.md anonymously', async () => { expect(requests[0]?.url).toContain('mode=rush') expect(requests[0]?.url).toContain('objective=compare+plans') expect(requests[0]?.method).toBe('GET') - expect(requests[0]?.headers).toEqual({ - accept: 'application/json', - }) + expect(requests[0]?.headers).toEqual( + expect.objectContaining({ + accept: 'application/json', + }), + ) expect(result).toEqual({ content: [{ type: 'text', text: '# Pricing' }], details: { @@ -721,11 +725,13 @@ test('uses session auth headers when available', async () => { const { tools } = loadExtension() await tools[0]!.execute('call_1', { url: 'https://example.com' }) - expect(requests[0]?.headers).toEqual({ - accept: 'application/json', - authorization: 'Bearer access-token-1', - 'x-organization-id': 'org_123', - }) + expect(requests[0]?.headers).toEqual( + expect.objectContaining({ + accept: 'application/json', + authorization: 'Bearer access-token-1', + 'x-organization-id': 'org_123', + }), + ) Session.delete() fs.rmSync(tmpDir, { force: true, recursive: true }) @@ -840,10 +846,12 @@ test('prefers CURLMD_API_KEY for authentication', async () => { await tools[0]!.execute('call_1', { url: 'https://example.com' }) expect(requests[0]?.method).toBe('GET') - expect(requests[0]?.headers).toEqual({ - accept: 'application/json', - authorization: 'Bearer curlmd_test_token', - }) + expect(requests[0]?.headers).toEqual( + expect.objectContaining({ + accept: 'application/json', + authorization: 'Bearer curlmd_test_token', + }), + ) }) test('throws validation issues for bad requests', async () => { diff --git a/plugins/pi/src/index.ts b/plugins/pi/src/index.ts index fe27d674..ac94e462 100644 --- a/plugins/pi/src/index.ts +++ b/plugins/pi/src/index.ts @@ -14,6 +14,8 @@ import { parseNumberHeader, } from './utils.ts' +const aiAgent = 'pi' as const + export default function (pi: ExtensionAPI) { const baseUrl = process.env.CURLMD_BASE_URL || defaultBaseUrl const apiKey = process.env.CURLMD_API_KEY @@ -123,7 +125,7 @@ export default function (pi: ExtensionAPI) { return } - const client = createClient(baseUrl, { headers: createHeaders(authHeaders) }) + const client = createClient(baseUrl, { aiAgent, headers: createHeaders(authHeaders) }) const [orgsRes, meRes] = await Promise.all([ client.api.orgs.$get(), client.api.auth.me.$get(), @@ -255,6 +257,7 @@ export default function (pi: ExtensionAPI) { const status = await (async () => { try { const client = createClient(baseUrl, { + aiAgent, headers: createHeaders({ authorization: authHeaders.authorization, expires_at: null, @@ -413,6 +416,7 @@ export default function (pi: ExtensionAPI) { })() const client = createClient(baseUrl, { + aiAgent, headers: apiKey ? createHeaders(null) : createHeaders(authHeaders), }) let res = await client.fetch(params.url, { @@ -428,6 +432,7 @@ export default function (pi: ExtensionAPI) { authHeaders = await resolver({ forceRefresh: true }) if (!authHeaders) authType = 'anon' const client = createClient(baseUrl, { + aiAgent, headers: apiKey ? createHeaders(null) : createHeaders(authHeaders), }) res = await client.fetch(params.url, { diff --git a/src/api.ts b/src/api.ts index 0e14d278..de0ee8da 100644 --- a/src/api.ts +++ b/src/api.ts @@ -11,6 +11,7 @@ import { estimateTokenCount } from 'tokenx' import { stringify as yamlStringify } from 'yaml' import { z } from 'zod' import { createClient, type Database } from '#db/client.ts' +import * as DbSchema from '#db/schemas.gen.ts' import type { DB } from '#db/types.gen.ts' import { requestTokensSavedSql } from '#db/utils.ts' import * as ApiKey from '#lib/apiKey.ts' @@ -1257,7 +1258,9 @@ export const api = new Hono<{ z.object({ expires_in: z.number().int().positive().optional(), max_uses: z.number().int().positive().nullable().default(null), - role: z.enum(['member', 'admin']).default('member'), + role: DbSchema.organization_invite.shape.role + .extract(['admin', 'member']) + .default('member'), }), ), async (c) => { @@ -1394,7 +1397,9 @@ export const api = new Hono<{ 'json', z.object({ login: z.string(), - role: z.enum(['member', 'admin']).default('member'), + role: DbSchema.organization_invite.shape.role + .extract(['admin', 'member']) + .default('member'), }), ), async (c) => { @@ -1446,7 +1451,10 @@ export const api = new Hono<{ ) .patch( '/api/orgs/:id/members/:memberId', - hono.validator('json', z.object({ role: z.enum(['member', 'admin']) })), + hono.validator( + 'json', + z.object({ role: DbSchema.organization_invite.shape.role.extract(['admin', 'member']) }), + ), async (c) => { if (hono.narrowValidation) return hono.validationError(c) if (!c.var.session) @@ -1839,7 +1847,7 @@ export const api = new Hono<{ .string() .transform((v) => v.split(/[\s,]+/).filter(Boolean)) .optional() - const mode = z.enum(['rush', 'smart']).default('smart') + const mode = DbSchema.request.shape.mode.unwrap().default('smart') const objective = z.string().optional() return z .object({ @@ -2228,8 +2236,17 @@ export const api = new Hono<{ })() const requestId = Nanoid.generate() + const aiAgent = (() => { + const parsed = z.safeParse( + DbSchema.request.shape.ai_agent.unwrap(), + c.req.header('x-ai-agent'), + ) + if (!parsed.success) return null + return parsed.data + })() await c.env.REQUEST_QUEUE.send({ account_id: c.var.session?.account_id ?? null, + ai_agent: aiAgent, api_key_id: c.var.api_key_id, billable, cached, diff --git a/src/og.tsx b/src/og.tsx index ea8b13f9..7f21b71a 100644 --- a/src/og.tsx +++ b/src/og.tsx @@ -28,8 +28,7 @@ async function getElement(env: Cloudflare.Env, db: Database, query: query) { return urlVariant(query.url, tokensSaved) } case 'playground': { - const tokensSaved = await getTokensSaved(env, db) - return playgroundVariant(tokensSaved) + return playgroundVariant() } case 'index': { const tokensSaved = await getTokensSaved(env, db) @@ -65,25 +64,19 @@ function indexVariant(tokensSaved: number) {
tokens saved
-
{tokensSaved.toLocaleString()}
+
{formatNumber(tokensSaved)}
) } -function playgroundVariant(tokensSaved: number) { +function playgroundVariant() { return (
playground
-
-
- tokens saved -
-
{tokensSaved.toLocaleString()}
-
) } @@ -101,7 +94,7 @@ function urlVariant(urlParam: string, tokensSaved: number) {
tokens saved
-
{tokensSaved.toLocaleString()}
+
{formatNumber(tokensSaved)}
) @@ -154,6 +147,10 @@ export async function render(request: Request, env: Cloudflare.Env, db: Database }) } +function formatNumber(n: number) { + return n.toString().replace(/\B(?=(\d{3})+(?!\d))/g, ',') +} + async function loadFont(request: Request, env: Cloudflare.Env, path: string) { const url = new URL(path, request.url) return env.ASSETS.fetch(url).then((r) => r.arrayBuffer()) diff --git a/src/queues/request.ts b/src/queues/request.ts index 824d3404..0c748f44 100644 --- a/src/queues/request.ts +++ b/src/queues/request.ts @@ -14,6 +14,7 @@ export async function processRequestMessage( .insertInto('request') .values({ account_id: body.account_id, + ai_agent: body.ai_agent ?? null, api_key_id: body.api_key_id, cached: body.cached, extracted_tokens: body.extracted_tokens, @@ -90,6 +91,7 @@ processRequestMessage.queueName = 'curl-request' as const export namespace processRequestMessage { export type Body = { account_id: string | null + ai_agent?: DB.request['ai_agent'] api_key_id: string | null billable: boolean cached: boolean diff --git a/src/queues/request.workers.test.ts b/src/queues/request.workers.test.ts index 90ca04c6..82287b02 100644 --- a/src/queues/request.workers.test.ts +++ b/src/queues/request.workers.test.ts @@ -20,6 +20,7 @@ test('inserts request record', async () => { attempts: 1, body: { account_id: null, + ai_agent: 'amp', api_key_id: null, billable: false, cached: false, @@ -50,6 +51,7 @@ test('inserts request record', async () => { .where('id', '=', 'req_1') .selectAll() .executeTakeFirstOrThrow() + expect(row.ai_agent).toBe('amp') expect(row.extracted_tokens).toBe(15) expect(row.filtered_tokens).toBe(20) expect(row.hostname).toBe('example.com') From 001b4b6f52d37d763571c2d3da3ad425494efc78 Mon Sep 17 00:00:00 2001 From: tmm Date: Sun, 26 Apr 2026 17:03:42 -0400 Subject: [PATCH 2/2] chore: changeset --- .changeset/tricky-spies-give.md | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 .changeset/tricky-spies-give.md diff --git a/.changeset/tricky-spies-give.md b/.changeset/tricky-spies-give.md new file mode 100644 index 00000000..53f98bb8 --- /dev/null +++ b/.changeset/tricky-spies-give.md @@ -0,0 +1,5 @@ +--- +'curl.md': patch +--- + +Added agent detection