diff --git a/README.md b/README.md index 20813de3eb3..e3fa22f8abf 100644 --- a/README.md +++ b/README.md @@ -25,6 +25,7 @@ Learn all about Gemini CLI in our [documentation](https://geminicli.com/docs/). - **🔌 Extensible**: MCP (Model Context Protocol) support for custom integrations. - **💻 Terminal-first**: Designed for developers who live in the command line. +- **⚡ SGLang Inference Server**: Direct connection to local or remote SGLang servers (Kimi-K3, DeepSeek, Qwen). See [SGLang Setup Guide](docs/sglang.md). - **🛡️ Open source**: Apache 2.0 licensed. ## 📦 Installation @@ -71,6 +72,29 @@ conda activate gemini_env npm install -g @google/gemini-cli ``` +#### Build from Source with SGLang Support (Linux / macOS) + +To connect Gemini CLI to local or remote SGLang / OpenAI inference servers (e.g., Moonshot Kimi-K3): + +```bash +# 1. Clone the branch +git clone -b feat/sglang-support https://github.com/shivajid/gemini-cli.git +cd gemini-cli + +# 2. Install dependencies & build binary bundle +npm install +npm run build +npm run bundle +npm link + +# 3. Configure and run with SGLang +export SGLANG_BASE_URL="http://127.0.0.1:30100/v1" +export GEMINI_MODEL="moonshotai/Kimi-K3" +gemini +``` + +See the [SGLang Setup Guide](docs/sglang.md) for full Linux prerequisites and troubleshooting. + ## Release Channels See [Releases](https://www.geminicli.com/docs/changelogs) for more details. diff --git a/docs/sglang.md b/docs/sglang.md new file mode 100644 index 00000000000..2bf16de34e0 --- /dev/null +++ b/docs/sglang.md @@ -0,0 +1,189 @@ +# Connecting Gemini CLI to SGLang Server + +Gemini CLI includes native support for connecting directly to local or remote **SGLang inference servers** (such as [Moonshot Kimi-K3](https://huggingface.co/moonshotai/Kimi-K3), DeepSeek-V3/R1, or Qwen models). + +This integration leverages OpenAI-compatible `/v1/chat/completions` endpoints with full support for: +- ⚡ **Streaming responses** (`stream: true`) +- 🧠 **Reasoning thought traces** (e.g. `delta.reasoning_content`) rendered cleanly in the CLI thinking box +- 🛠️ **Built-in tools & MCP function calling** with recursive schema conversion (handling Gemini uppercase types to standard JSON Schema) +- 🛑 **Interactive stream cancellation** (`ESC` key support) +- 🔁 **Multi-turn conversation history** with persistent tool call identifiers + +--- + +## 1. Prerequisites on Linux (Debian / Ubuntu / COS) + +Before building, ensure you have **Node.js 20+**, **npm**, and build essentials installed on your Linux machine: + +### Install Node.js 20 and Build Tools + +```bash +# Update package lists +sudo apt-get update + +# Install git, curl, and native compilation tools +sudo apt-get install -y git curl build-essential python3 + +# Install Node.js 20.x LTS via NodeSource +curl -fsSL https://deb.nodesource.com/setup_20.x | sudo -E bash - +sudo apt-get install -y nodejs + +# Verify versions +node -v # Should be v20.x or higher +npm -v # Should be 10.x or higher +``` + +### Install GKE / Kubernetes Client Tools (If running on GKE) + +```bash +# Install kubectl and GKE auth plugin +sudo apt-get install -y kubectl google-cloud-cli-gke-gcloud-auth-plugin + +# Configure cluster credentials +gcloud container clusters get-credentials \ + --region \ + --project +``` + +--- + +## 2. Building `feat/sglang-support` from Source + +Clone the repository, switch to the `feat/sglang-support` branch, install dependencies, and compile: + +```bash +# 1. Clone the repository and checkout the feat/sglang-support branch +git clone -b feat/sglang-support https://github.com/shivajid/gemini-cli.git +cd gemini-cli + +# 2. Install workspace dependencies +npm install + +# 3. Compile all packages (including @google/gemini-cli-core) +npm run build + +# 4. Generate the standalone CLI binary bundle +npm run bundle + +# 5. Link globally so the `gemini` command is available system-wide +npm link +``` + +> **Tip for updating**: To pull future updates, simply run: +> ```bash +> git pull origin feat/sglang-support +> npm run build && npm run bundle && npm link +> ``` + +--- + +## 3. Quick Start & Connecting to SGLang + +### Step 1: Start SGLang Port-Forwarding + +If your SGLang server is running in Kubernetes / GKE, forward the API port to your local machine: + +```bash +# Forward port 30100 from your SGLang leader pod +kubectl port-forward -n pod/ 30100:30100 & +``` + +Verify that the server is reachable: + +```bash +curl http://127.0.0.1:30100/v1/models +``` + +--- + +### Step 2: Set Environment Variables + +```bash +export SGLANG_BASE_URL="http://127.0.0.1:30100/v1" +export GEMINI_MODEL="moonshotai/Kimi-K3" +export GEMINI_DEFAULT_AUTH_TYPE="sglang" +``` + +> **Note**: Always use `http://127.0.0.1:30100/v1` instead of `localhost` on Linux containers to avoid DNS resolution issues. + +--- + +### Step 3: Configure Settings (Optional) + +Create or update `~/.gemini/settings.json` to persist the SGLang authentication: + +```bash +mkdir -p ~/.gemini +cat << 'EOF' > ~/.gemini/settings.json +{ + "general": { + "enableAutoUpdateNotification": false + }, + "security": { + "auth": { + "selectedType": "sglang" + } + } +} +EOF +``` + +--- + +### Step 4: Run Gemini CLI + +Start an interactive chat session: + +```bash +gemini +``` + +Or pass an immediate prompt: + +```bash +gemini "Hello Kimi-K3! List the files in this directory." +``` + +--- + +## 4. Interactive Authentication Menu + +If you run `gemini` without predefined settings, or type `/auth` inside an active session: + +``` +? How would you like to authenticate for this project? +● 1. SGLang Server (Local / Remote Kimi-K3) + 2. Sign in with Google + 3. Use Gemini API Key + 4. Vertex AI +``` + +Select **`1. SGLang Server (Local / Remote Kimi-K3)`** to bypass Google credentials and route traffic directly to your SGLang endpoint. + +--- + +## 5. Configuration Reference + +| Variable / Setting | Description | Default | +|---|---|---| +| `SGLANG_BASE_URL` | Base URL of the OpenAI-compatible SGLang server | `http://127.0.0.1:30100/v1` | +| `OPENAI_BASE_URL` | Secondary fallback base URL | `http://127.0.0.1:30100/v1` | +| `GEMINI_MODEL` / `SGLANG_MODEL` | Served model name on SGLang | `moonshotai/Kimi-K3` | +| `GEMINI_DEFAULT_AUTH_TYPE` | Default auth method (`sglang`, `oauth-personal`, `gemini-api-key`) | `oauth-personal` | +| `enableAutoUpdateNotification` | Set to `false` in `settings.json` to hide git update banners | `true` | + +--- + +## 6. Troubleshooting + +### 1. `socket.gaierror: [Errno -2] Name or service not known` +- **Cause**: Linux environment does not resolve `localhost` in `/etc/hosts`. +- **Fix**: Use numeric IP `http://127.0.0.1:30100/v1`. + +### 2. `Connection refused` +- **Cause**: SGLang server is initializing or `kubectl port-forward` terminated. +- **Fix**: Check `kubectl get pods -n ` and restart port-forwarding. + +### 3. `Model "moonshotai/Kimi-K3" was not found` +- **Cause**: Saved setting in `~/.gemini/settings.json` is still set to Google API (`gemini-api-key` or `oauth-personal`). +- **Fix**: Run `/auth` and select **SGLang Server**, or set `"selectedType": "sglang"` in `~/.gemini/settings.json`. diff --git a/packages/cli/src/config/auth.ts b/packages/cli/src/config/auth.ts index 1ca07f98eb4..cdf9b730cdd 100644 --- a/packages/cli/src/config/auth.ts +++ b/packages/cli/src/config/auth.ts @@ -13,7 +13,8 @@ export async function validateAuthMethod( loadEnvironment(loadSettings().merged, process.cwd()); if ( authMethod === AuthType.LOGIN_WITH_GOOGLE || - authMethod === AuthType.COMPUTE_ADC + authMethod === AuthType.COMPUTE_ADC || + authMethod === AuthType.SGLANG ) { return null; } diff --git a/packages/cli/src/ui/auth/AuthDialog.tsx b/packages/cli/src/ui/auth/AuthDialog.tsx index e5b08e6cc2d..e0c046a667a 100644 --- a/packages/cli/src/ui/auth/AuthDialog.tsx +++ b/packages/cli/src/ui/auth/AuthDialog.tsx @@ -43,6 +43,11 @@ export function AuthDialog({ }: AuthDialogProps): React.JSX.Element { const [exiting, setExiting] = useState(false); let items = [ + { + label: 'SGLang Server (Local / Remote Kimi-K3)', + value: AuthType.SGLANG, + key: AuthType.SGLANG, + }, { label: 'Sign in with Google', value: AuthType.LOGIN_WITH_GOOGLE, diff --git a/packages/cli/src/ui/auth/useAuth.ts b/packages/cli/src/ui/auth/useAuth.ts index caa9ed2c4b8..5107c196113 100644 --- a/packages/cli/src/ui/auth/useAuth.ts +++ b/packages/cli/src/ui/auth/useAuth.ts @@ -29,8 +29,12 @@ export async function validateAuthMethodWithSettings( if (settings.merged.security.auth.useExternal) { return null; } - // If using Gemini API key, we don't validate it here as we might need to prompt for it. - if (authType === AuthType.USE_GEMINI) { + // If using Gemini API key or SGLang, we don't validate it here as we might need to prompt for it. + if ( + authType === AuthType.USE_GEMINI || + authType === AuthType.SGLANG || + (authType as string) === 'sglang' + ) { return null; } return validateAuthMethod(authType); diff --git a/packages/core/src/config/models.ts b/packages/core/src/config/models.ts index dd6506b837c..8d8c74a95da 100644 --- a/packages/core/src/config/models.ts +++ b/packages/core/src/config/models.ts @@ -84,6 +84,7 @@ export const PREVIEW_GEMINI_FLASH_LITE_MODEL = 'none'; export const GEMMA_4_31B_IT_MODEL = 'gemma-4-31b-it'; export const GEMMA_4_26B_A4B_IT_MODEL = 'gemma-4-26b-a4b-it'; +export const KIMI_K3_MODEL = 'moonshotai/Kimi-K3'; export const VALID_GEMINI_MODELS = new Set([ PREVIEW_GEMINI_MODEL, @@ -99,6 +100,7 @@ export const VALID_GEMINI_MODELS = new Set([ GEMMA_4_31B_IT_MODEL, GEMMA_4_26B_A4B_IT_MODEL, + KIMI_K3_MODEL, ]); /** @deprecated Use GEMINI_MODEL_ALIAS_AUTO instead. */ @@ -549,6 +551,9 @@ export function isActiveModel( useCustomToolModel: boolean = false, experimentalGemma: boolean = true, ): boolean { + if (model === KIMI_K3_MODEL || model.startsWith('moonshotai/') || model.includes('kimi')) { + return true; + } if (!VALID_GEMINI_MODELS.has(model) || model === 'none') { return false; } diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts index e0b33d8c7b0..48a5f560c3a 100644 --- a/packages/core/src/core/contentGenerator.ts +++ b/packages/core/src/core/contentGenerator.ts @@ -32,6 +32,7 @@ import { getVersion, resolveModel } from '../../index.js'; import type { LlmRole } from '../telemetry/llmRole.js'; import { ModelMappingContentGenerator } from './modelMappingContentGenerator.js'; import { CCPA_AI_MODEL_MAPPINGS } from '../config/models.js'; +import { SglangContentGenerator } from './sglangContentGenerator.js'; /** * Interface abstracting the core functionalities for generating content and counting tokens. @@ -67,17 +68,22 @@ export enum AuthType { LEGACY_CLOUD_SHELL = 'cloud-shell', COMPUTE_ADC = 'compute-default-credentials', GATEWAY = 'gateway', + SGLANG = 'sglang', } /** * Detects the best authentication type based on environment variables. * * Checks in order: - * 1. GOOGLE_GENAI_USE_GCA=true -> LOGIN_WITH_GOOGLE - * 2. GOOGLE_GENAI_USE_VERTEXAI=true -> USE_VERTEX_AI - * 3. GEMINI_API_KEY -> USE_GEMINI + * 1. SGLANG_BASE_URL or OPENAI_BASE_URL -> SGLANG + * 2. GOOGLE_GENAI_USE_GCA=true -> LOGIN_WITH_GOOGLE + * 3. GOOGLE_GENAI_USE_VERTEXAI=true -> USE_VERTEX_AI + * 4. GEMINI_API_KEY -> USE_GEMINI */ export function getAuthTypeFromEnv(): AuthType | undefined { + if (process.env['SGLANG_BASE_URL'] || process.env['OPENAI_BASE_URL']) { + return AuthType.SGLANG; + } if (process.env['GOOGLE_GENAI_USE_GCA'] === 'true') { return AuthType.LOGIN_WITH_GOOGLE; } @@ -166,7 +172,8 @@ export async function createContentGeneratorConfig( // (WSL/SSH/Docker/CI) keytar can block indefinitely on its functional probe. if ( authType === AuthType.LOGIN_WITH_GOOGLE || - authType === AuthType.COMPUTE_ADC + authType === AuthType.COMPUTE_ADC || + authType === AuthType.SGLANG ) { return contentGeneratorConfig; } @@ -409,6 +416,36 @@ export async function createContentGenerator( }); return new LoggingContentGenerator(googleGenAI.models, gcConfig); } + if ( + config.authType === AuthType.SGLANG || + String(config.authType).toLowerCase() === 'sglang' + ) { + const baseUrl = + config.baseUrl || + process.env['SGLANG_BASE_URL'] || + process.env['OPENAI_BASE_URL'] || + 'http://127.0.0.1:30100/v1'; + validateBaseUrl(baseUrl); + // Prefer an explicitly configured served-model name; never forward + // internal gemini-* aliases to the SGLang server. + const configuredModel = ( + gcConfig as { getModel?: () => string } + ).getModel?.(); + const isServedModelName = (m?: string): m is string => + !!m && + m !== 'auto' && + m !== 'none' && + !m.startsWith('gemini') && + !m.startsWith('gemma'); + const modelName = + process.env['SGLANG_MODEL'] || + (isServedModelName(configuredModel) ? configuredModel : undefined) || + 'moonshotai/Kimi-K3'; + return new LoggingContentGenerator( + new SglangContentGenerator(baseUrl, modelName), + gcConfig, + ); + } throw new Error( `Error creating contentGenerator: Unsupported authType: ${config.authType}`, ); diff --git a/packages/core/src/core/sglangContentGenerator.test.ts b/packages/core/src/core/sglangContentGenerator.test.ts new file mode 100644 index 00000000000..6100cf22dda --- /dev/null +++ b/packages/core/src/core/sglangContentGenerator.test.ts @@ -0,0 +1,284 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; +import { GenerateContentResponse } from '@google/genai'; +import type { GenerateContentParameters } from '@google/genai'; +import { SglangContentGenerator } from './sglangContentGenerator.js'; +import { LlmRole } from '../telemetry/types.js'; + +function sseBody(chunks: unknown[]): ReadableStream { + const encoder = new TextEncoder(); + return new ReadableStream({ + start(controller) { + for (const chunk of chunks) { + controller.enqueue( + encoder.encode(`data: ${JSON.stringify(chunk)}\n\n`), + ); + } + controller.enqueue(encoder.encode('data: [DONE]\n\n')); + controller.close(); + }, + }); +} + +describe('SglangContentGenerator', () => { + const fetchMock = vi.fn(); + + beforeEach(() => { + vi.stubGlobal('fetch', fetchMock); + fetchMock.mockReset(); + }); + + afterEach(() => { + vi.unstubAllGlobals(); + }); + + const generator = () => + new SglangContentGenerator('http://localhost:9999/v1', 'moonshotai/Kimi-K3'); + + it('streams tool calls as real GenerateContentResponse instances with working functionCalls getter', async () => { + fetchMock.mockResolvedValue({ + ok: true, + body: sseBody([ + { + choices: [ + { + index: 0, + delta: { reasoning_content: 'thinking...' }, + finish_reason: null, + }, + ], + }, + { + choices: [ + { + index: 0, + delta: { + tool_calls: [ + { + index: 0, + id: 'read_file:0', + function: { name: 'read_file', arguments: '{"file_' }, + }, + ], + }, + finish_reason: null, + }, + ], + }, + { + choices: [ + { + index: 0, + delta: { + tool_calls: [ + { index: 0, function: { arguments: 'path":"/tmp/x"}' } }, + ], + }, + finish_reason: null, + }, + ], + }, + { choices: [{ index: 0, delta: {}, finish_reason: 'tool_calls' }] }, + { + choices: [], + usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 }, + }, + ]), + }); + + const stream = await generator().generateContentStream( + { + model: 'gemini-2.5-pro', + contents: 'hi', + } as GenerateContentParameters, + 'prompt-1', + LlmRole.MAIN, + ); + + const chunks: GenerateContentResponse[] = []; + for await (const chunk of stream) { + chunks.push(chunk); + } + + // Every chunk must be a real SDK class instance so getters work. + for (const chunk of chunks) { + expect(chunk).toBeInstanceOf(GenerateContentResponse); + } + + // Thought chunk is prefixed with a stable subject. + const thoughtChunk = chunks.find((c) => + c.candidates?.[0]?.content?.parts?.some((p) => p.thought), + ); + expect( + thoughtChunk?.candidates?.[0]?.content?.parts?.[0]?.text, + ).toContain('**Thinking**'); + + // Split tool-call arguments are reassembled, and the SDK + // `functionCalls` getter (used by turn.ts) sees the call. + const toolChunk = chunks.find( + (c) => c.functionCalls && c.functionCalls.length > 0, + ); + expect(toolChunk).toBeDefined(); + expect(toolChunk!.functionCalls![0]).toEqual({ + id: 'read_file:0', + name: 'read_file', + args: { file_path: '/tmp/x' }, + }); + // The finish chunk carries the finish reason. + expect(toolChunk!.candidates![0].finishReason).toBe('STOP'); + + // Usage from stream_options.include_usage is surfaced. + const usageChunk = chunks.find((c) => c.usageMetadata); + expect(usageChunk?.usageMetadata?.totalTokenCount).toBe(15); + + // The request advertised stream_options and mapped the gemini model + // alias to the served model name. + const payload = JSON.parse(fetchMock.mock.calls[0][1].body as string); + expect(payload.model).toBe('moonshotai/Kimi-K3'); + expect(payload.stream_options).toEqual({ include_usage: true }); + }); + + it('round-trips tool call ids and skips thought parts in history', async () => { + fetchMock.mockResolvedValue({ + ok: true, + json: async () => ({ + choices: [ + { message: { content: 'done' }, finish_reason: 'stop' }, + ], + usage: {}, + }), + }); + + await generator().generateContent( + { + model: 'moonshotai/Kimi-K3', + contents: [ + { role: 'user', parts: [{ text: 'read x' }] }, + { + role: 'model', + parts: [ + { text: 'internal reasoning', thought: true }, + { + functionCall: { + id: 'read_file__synth_1', + name: 'read_file', + args: { file_path: '/tmp/x' }, + }, + }, + ], + }, + { + role: 'user', + parts: [ + { + functionResponse: { + id: 'read_file__synth_1', + name: 'read_file', + response: { output: 'file contents' }, + }, + }, + ], + }, + ], + } as GenerateContentParameters, + 'prompt-1', + LlmRole.MAIN, + ); + + const payload = JSON.parse(fetchMock.mock.calls[0][1].body as string); + const messages = payload.messages as Array>; + + const assistant = messages.find((m) => m['role'] === 'assistant'); + const tool = messages.find((m) => m['role'] === 'tool'); + expect(assistant).toBeDefined(); + expect(tool).toBeDefined(); + + const toolCalls = assistant!['tool_calls'] as Array<{ + id: string; + function: { name: string; arguments: string }; + }>; + expect(toolCalls[0].id).toBe('read_file__synth_1'); + expect(tool!['tool_call_id']).toBe('read_file__synth_1'); + expect(tool!['content']).toBe('file contents'); + + // Reasoning must never be echoed back to the server. + expect(JSON.stringify(messages)).not.toContain('internal reasoning'); + }); + + it('converts parametersJsonSchema tool declarations', async () => { + fetchMock.mockResolvedValue({ + ok: true, + json: async () => ({ + choices: [{ message: { content: 'ok' }, finish_reason: 'stop' }], + }), + }); + + await generator().generateContent( + { + model: 'moonshotai/Kimi-K3', + contents: 'hi', + config: { + tools: [ + { + functionDeclarations: [ + { + name: 'read_file', + description: 'Reads a file', + parametersJsonSchema: { + type: 'object', + properties: { file_path: { type: 'STRING' } }, + required: ['file_path'], + }, + }, + ], + }, + ], + }, + } as GenerateContentParameters, + 'prompt-1', + LlmRole.MAIN, + ); + + const payload = JSON.parse(fetchMock.mock.calls[0][1].body as string); + expect(payload.tools).toHaveLength(1); + expect(payload.tools[0].function.name).toBe('read_file'); + // Gemini upper-case types are normalized for OpenAI-compatible servers. + expect( + payload.tools[0].function.parameters.properties.file_path.type, + ).toBe('string'); + expect(payload.tools[0].function.parameters.required).toEqual([ + 'file_path', + ]); + }); + + it('honors abortSignal', async () => { + const controller = new AbortController(); + fetchMock.mockImplementation( + (_url: string, init: { signal?: AbortSignal }) => + new Promise((_resolve, reject) => { + init.signal?.addEventListener('abort', () => { + const err = new Error('aborted'); + err.name = 'AbortError'; + reject(err); + }); + }), + ); + + const promise = generator().generateContent( + { + model: 'moonshotai/Kimi-K3', + contents: 'hi', + config: { abortSignal: controller.signal }, + } as GenerateContentParameters, + 'prompt-1', + LlmRole.MAIN, + ); + controller.abort(); + await expect(promise).rejects.toThrow('aborted'); + }); +}); diff --git a/packages/core/src/core/sglangContentGenerator.ts b/packages/core/src/core/sglangContentGenerator.ts new file mode 100644 index 00000000000..db0a0b8c687 --- /dev/null +++ b/packages/core/src/core/sglangContentGenerator.ts @@ -0,0 +1,772 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/* eslint-disable @typescript-eslint/no-unsafe-type-assertion, no-restricted-syntax -- + * This file converts between the Gemini SDK types and the untyped + * OpenAI-compatible wire format served by SGLang; structural checks and + * assertions against parsed JSON are unavoidable here. */ + +import { + GenerateContentResponse, + FinishReason, + type CountTokensParameters, + type CountTokensResponse, + type EmbedContentParameters, + type EmbedContentResponse, + type GenerateContentParameters, + type GenerateContentConfig, + type Content, + type ContentUnion, + type Part, + type FunctionDeclaration, +} from '@google/genai'; +import type { ContentGenerator } from './contentGenerator.js'; +import type { LlmRole } from '../telemetry/types.js'; + +interface OpenAIToolCall { + id: string; + type: 'function'; + function: { name: string; arguments: string }; +} + +interface OpenAIMessage { + role: 'system' | 'user' | 'assistant' | 'tool'; + content?: string | null; + reasoning_content?: string; + tool_calls?: OpenAIToolCall[]; + tool_call_id?: string; +} + +interface OpenAIUsage { + prompt_tokens?: number; + completion_tokens?: number; + total_tokens?: number; +} + +/** + * Recursively normalizes a JSON schema for OpenAI-compatible servers: + * lowercases `type` values (Gemini uses upper case enums) and recurses + * into all nested schema locations. + */ +function cleanSchema(schema: unknown): unknown { + if (!schema || typeof schema !== 'object') return schema; + if (Array.isArray(schema)) { + return schema.map(cleanSchema); + } + const obj = { ...(schema as Record) }; + delete obj['$schema']; + if (typeof obj['type'] === 'string') { + obj['type'] = obj['type'].toLowerCase(); + } else if (Array.isArray(obj['type'])) { + obj['type'] = (obj['type'] as unknown[]).map((t) => + typeof t === 'string' ? t.toLowerCase() : t, + ); + } + for (const key of ['properties', '$defs', 'definitions']) { + if (obj[key] && typeof obj[key] === 'object' && !Array.isArray(obj[key])) { + const mapped: Record = {}; + for (const [k, v] of Object.entries( + obj[key] as Record, + )) { + mapped[k] = cleanSchema(v); + } + obj[key] = mapped; + } + } + for (const key of [ + 'items', + 'additionalProperties', + 'anyOf', + 'oneOf', + 'allOf', + 'prefixItems', + 'not', + ]) { + if (obj[key] !== undefined && typeof obj[key] === 'object') { + obj[key] = cleanSchema(obj[key]); + } + } + return obj; +} + +function mapFinishReason(reason?: string | null): FinishReason | undefined { + if (!reason) return undefined; + switch (reason) { + case 'length': + return FinishReason.MAX_TOKENS; + case 'content_filter': + return FinishReason.SAFETY; + case 'stop': + case 'tool_calls': + case 'function_call': + default: + return FinishReason.STOP; + } +} + +function extractText(content: ContentUnion | string | undefined): string { + if (!content) return ''; + if (typeof content === 'string') return content; + if (Array.isArray(content)) { + return content + .map((p) => (typeof p === 'string' ? p : (p.text ?? ''))) + .filter(Boolean) + .join('\n'); + } + const maybeContent = content as Content; + if (Array.isArray(maybeContent.parts)) { + return maybeContent.parts + .map((p) => p.text) + .filter(Boolean) + .join('\n'); + } + return (content as Part).text ?? ''; +} + +/** + * Extracts a plain-text payload from a Gemini functionResponse to send as + * an OpenAI `tool` message. + */ +function functionResponseToText(response: unknown): string { + if (response === null || response === undefined) return ''; + if (typeof response === 'string') return response; + if (typeof response === 'object') { + const obj = response as Record; + if (typeof obj['output'] === 'string') return obj['output']; + if (obj['error'] !== undefined) { + return typeof obj['error'] === 'string' + ? `Error: ${obj['error']}` + : `Error: ${JSON.stringify(obj['error'])}`; + } + } + try { + return JSON.stringify(response); + } catch { + return String(response); + } +} + +/** + * ContentGenerator implementation for an SGLang (OpenAI-compatible) server, + * e.g. a multi-node Kimi-K3 deployment launched with + * `--tool-call-parser kimi_k3 --reasoning-parser kimi_k3`. + * + * Notes on correctness (things the rest of gemini-cli relies on): + * - Responses MUST be real `GenerateContentResponse` instances: turn.ts and + * geminiChat.ts use the SDK class getters (`resp.functionCalls`) to + * dispatch tool calls. Plain object literals silently break tool use. + * - functionCall parts must carry stable `id`s and history conversion must + * reuse them so `tool_calls[].id` matches the `tool` message + * `tool_call_id` across turns. + * - `config.abortSignal` must be honored so ESC cancels generation. + * - Never write to stdout/stderr: it corrupts the Ink UI. + */ +export class SglangContentGenerator implements ContentGenerator { + private baseUrl: string; + private defaultModel: string; + private apiKey?: string; + private syntheticIdCounter = 0; + + constructor( + baseUrl: string = 'http://127.0.0.1:30100/v1', + defaultModel: string = 'moonshotai/Kimi-K3', + apiKey?: string, + ) { + this.baseUrl = baseUrl.replace(/\/+$/, ''); + this.defaultModel = defaultModel; + this.apiKey = + apiKey || + process.env['SGLANG_API_KEY'] || + process.env['OPENAI_API_KEY'] || + undefined; + } + + /** + * The CLI resolves gemini-* model aliases internally; the sglang server + * only knows its own served model name. Map anything that doesn't look + * like a served model onto the configured default. + */ + private resolveModelName(requested?: string | null): string { + if (!requested) return this.defaultModel; + if ( + requested === 'auto' || + requested.startsWith('gemini') || + requested.startsWith('gemma') || + requested === 'none' + ) { + return this.defaultModel; + } + return requested; + } + + private normalizeContents( + contents: GenerateContentParameters['contents'], + ): Content[] { + if (!contents) return []; + if (typeof contents === 'string') { + return [{ role: 'user', parts: [{ text: contents }] }]; + } + const list = Array.isArray(contents) ? contents : [contents]; + const result: Content[] = []; + for (const item of list) { + if (typeof item === 'string') { + result.push({ role: 'user', parts: [{ text: item }] }); + } else if ( + item && + typeof item === 'object' && + ('role' in item || 'parts' in item) + ) { + result.push(item); + } else if (item && typeof item === 'object') { + // A bare Part. + result.push({ role: 'user', parts: [item as Part] }); + } + } + return result; + } + + private convertContentsToMessages( + contents: GenerateContentParameters['contents'], + systemInstruction?: ContentUnion | string, + ): OpenAIMessage[] { + const messages: OpenAIMessage[] = []; + + const systemText = extractText(systemInstruction); + if (systemText) { + messages.push({ role: 'system', content: systemText }); + } + + // Queue of auto-assigned tool call ids, used to pair functionResponses + // that arrive without an id to the calls they answer (FIFO order). + const pendingAutoIds: string[] = []; + + for (const c of this.normalizeContents(contents)) { + if (!c.parts || c.parts.length === 0) continue; + const isModel = c.role === 'model'; + + let text = ''; + const toolCalls: OpenAIToolCall[] = []; + const toolResponses: Array<{ id: string; content: string }> = []; + + for (const part of c.parts) { + // Never re-send reasoning/thought parts back to the server. + if (part.thought) continue; + + if (part.text) { + text += (text ? '\n' : '') + part.text; + } + if (part.functionCall) { + let id = part.functionCall.id; + if (!id) { + id = `call_auto_${this.syntheticIdCounter++}`; + pendingAutoIds.push(id); + } + toolCalls.push({ + id, + type: 'function', + function: { + name: part.functionCall.name || 'unknown_tool', + arguments: JSON.stringify(part.functionCall.args ?? {}), + }, + }); + } + if (part.functionResponse) { + const id = + part.functionResponse.id || + pendingAutoIds.shift() || + `call_auto_${this.syntheticIdCounter++}`; + toolResponses.push({ + id, + content: functionResponseToText(part.functionResponse.response), + }); + } + } + + // Tool results must directly follow the assistant message that + // requested them. + for (const tr of toolResponses) { + messages.push({ + role: 'tool', + tool_call_id: tr.id, + content: tr.content, + }); + } + + if (toolCalls.length > 0) { + messages.push({ + role: 'assistant', + content: text || null, + tool_calls: toolCalls, + }); + } else if (text) { + messages.push({ + role: isModel ? 'assistant' : 'user', + content: text, + }); + } + } + + if (messages.filter((m) => m.role !== 'system').length === 0) { + messages.push({ role: 'user', content: 'Hello' }); + } + + return messages; + } + + private convertTools(config?: GenerateContentConfig) { + if (!config?.tools || !Array.isArray(config.tools)) { + return undefined; + } + const tools: Array<{ + type: 'function'; + function: { name: string; description?: string; parameters?: unknown }; + }> = []; + + for (const t of config.tools) { + const declarations = ( + t as { functionDeclarations?: FunctionDeclaration[] } + ).functionDeclarations; + if (!Array.isArray(declarations)) continue; + for (const fd of declarations) { + if (!fd.name) continue; + // gemini-cli built-in tools declare `parametersJsonSchema` + // (raw JSON schema); MCP/legacy tools may use `parameters`. + const schema = + (fd as { parametersJsonSchema?: unknown }).parametersJsonSchema ?? + fd.parameters; + tools.push({ + type: 'function', + function: { + name: fd.name, + description: fd.description, + parameters: cleanSchema(schema) ?? { + type: 'object', + properties: {}, + }, + }, + }); + } + } + return tools.length > 0 ? tools : undefined; + } + + private buildPayload( + request: GenerateContentParameters, + stream: boolean, + ): Record { + const config = request.config; + const messages = this.convertContentsToMessages( + request.contents, + config?.systemInstruction, + ); + const tools = this.convertTools(config); + + const payload: Record = { + model: this.resolveModelName(request.model), + messages, + stream, + // Large default: Kimi-K3 edits/writes whole files through tool + // arguments; a small cap truncates them into malformed JSON. + max_tokens: config?.maxOutputTokens ?? 32768, + }; + if (stream) { + payload['stream_options'] = { include_usage: true }; + } + if (config?.temperature !== undefined) { + payload['temperature'] = config.temperature; + } + if (config?.topP !== undefined) { + payload['top_p'] = config.topP; + } + if (config?.stopSequences && config.stopSequences.length > 0) { + payload['stop'] = config.stopSequences; + } + if (tools) { + payload['tools'] = tools; + const mode = config?.toolConfig?.functionCallingConfig?.mode; + if (mode === 'ANY') { + payload['tool_choice'] = 'required'; + } else if (mode === 'NONE') { + payload['tool_choice'] = 'none'; + } + } + + // Structured output for internal utility calls (generateJson): + // next-speaker checks, loop detection, summarization, etc. + const jsonSchema = + (config as { responseJsonSchema?: unknown })?.responseJsonSchema ?? + config?.responseSchema; + if (jsonSchema) { + payload['response_format'] = { + type: 'json_schema', + json_schema: { + name: 'response', + schema: cleanSchema(jsonSchema), + }, + }; + } else if (config?.responseMimeType === 'application/json') { + payload['response_format'] = { type: 'json_object' }; + } + + return payload; + } + + private async postChatCompletions( + payload: Record, + abortSignal?: AbortSignal, + ): Promise { + const headers: Record = { + 'Content-Type': 'application/json', + }; + if (this.apiKey) { + headers['Authorization'] = `Bearer ${this.apiKey}`; + } + let res: Response; + try { + res = await fetch(`${this.baseUrl}/chat/completions`, { + method: 'POST', + headers, + body: JSON.stringify(payload), + signal: abortSignal ?? null, + }); + } catch (err) { + if (err instanceof Error && err.name === 'AbortError') { + throw err; + } + throw new Error( + `Failed to reach SGLang server at ${this.baseUrl}/chat/completions. ` + + `Is kubectl port-forward running? ` + + `(${err instanceof Error ? err.message : String(err)})`, + ); + } + if (!res.ok) { + const errText = await res.text().catch(() => ''); + throw new Error(`SGLang server error (${res.status}): ${errText}`); + } + return res; + } + + private makeResponse( + parts: Part[], + finishReason?: FinishReason, + usage?: OpenAIUsage, + responseId?: string, + modelVersion?: string, + ): GenerateContentResponse { + const response = new GenerateContentResponse(); + response.candidates = [ + { + content: { parts, role: 'model' }, + ...(finishReason ? { finishReason } : {}), + index: 0, + }, + ]; + if (usage) { + response.usageMetadata = { + promptTokenCount: usage.prompt_tokens ?? 0, + candidatesTokenCount: usage.completion_tokens ?? 0, + totalTokenCount: usage.total_tokens ?? 0, + }; + } + if (responseId) { + response.responseId = responseId; + } + if (modelVersion) { + response.modelVersion = modelVersion; + } + return response; + } + + async generateContent( + request: GenerateContentParameters, + _userPromptId: string, + _role: LlmRole, + ): Promise { + const payload = this.buildPayload(request, false); + const res = await this.postChatCompletions( + payload, + request.config?.abortSignal, + ); + const data = (await res.json()) as { + id?: string; + model?: string; + choices?: Array<{ + message?: { + content?: string | null; + reasoning_content?: string | null; + tool_calls?: Array<{ + id?: string; + function?: { name?: string; arguments?: string }; + }>; + }; + finish_reason?: string; + }>; + usage?: OpenAIUsage; + }; + + const choice = data.choices?.[0]; + const message = choice?.message; + const parts: Part[] = []; + + if (message?.reasoning_content) { + parts.push({ text: message.reasoning_content, thought: true }); + } + if (message?.content) { + parts.push({ text: message.content }); + } + if (Array.isArray(message?.tool_calls)) { + for (const tc of message.tool_calls) { + if (!tc.function?.name) continue; + let args: Record = {}; + try { + args = JSON.parse(tc.function.arguments || '{}') as Record< + string, + unknown + >; + } catch { + // Leave args empty if the server returned malformed JSON. + } + parts.push({ + functionCall: { + id: tc.id || `call_auto_${this.syntheticIdCounter++}`, + name: tc.function.name, + args, + }, + }); + } + } + + return this.makeResponse( + parts, + mapFinishReason(choice?.finish_reason) ?? FinishReason.STOP, + data.usage, + data.id, + data.model, + ); + } + + async generateContentStream( + request: GenerateContentParameters, + _userPromptId: string, + _role: LlmRole, + ): Promise> { + const payload = this.buildPayload(request, true); + const res = await this.postChatCompletions( + payload, + request.config?.abortSignal, + ); + if (!res.body) { + throw new Error('SGLang server returned an empty stream body.'); + } + + const reader = res.body.getReader(); + const decoder = new TextDecoder(); + // eslint-disable-next-line @typescript-eslint/no-this-alias + const self = this; + + async function* makeStream(): AsyncGenerator { + let buffer = ''; + let sentFinish = false; + let usage: OpenAIUsage | undefined; + let responseId: string | undefined; + let modelVersion: string | undefined; + const toolCalls = new Map< + number, + { id?: string; name?: string; arguments: string } + >(); + + const flushToolCallParts = (): Part[] => { + const parts: Part[] = []; + const sorted = [...toolCalls.entries()].sort((a, b) => a[0] - b[0]); + for (const [, tc] of sorted) { + if (!tc.name) continue; + let args: Record = {}; + try { + args = JSON.parse(tc.arguments || '{}') as Record; + } catch { + // Malformed JSON args; send empty args rather than crashing. + } + parts.push({ + functionCall: { + id: tc.id || `call_auto_${self.syntheticIdCounter++}`, + name: tc.name, + args, + }, + }); + } + toolCalls.clear(); + return parts; + }; + + const handleData = function* ( + json: Record, + ): Generator { + if (typeof json['id'] === 'string') responseId = json['id']; + if (typeof json['model'] === 'string') modelVersion = json['model']; + if (json['usage']) { + usage = json['usage'] as OpenAIUsage; + } + const choice = ( + json['choices'] as + | Array<{ + delta?: { + content?: string | null; + reasoning_content?: string | null; + tool_calls?: Array<{ + index?: number; + id?: string; + function?: { name?: string; arguments?: string }; + }>; + }; + finish_reason?: string | null; + }> + | undefined + )?.[0]; + if (!choice) return; + + const delta = choice.delta; + const parts: Part[] = []; + + if (delta?.reasoning_content) { + // Prefix with a bold subject so gemini-cli's thought parser + // (which expects `**Subject** description`) renders a stable + // "Thinking" label instead of arbitrary reasoning fragments. + parts.push({ + text: `**Thinking** ${delta.reasoning_content}`, + thought: true, + }); + } + if (delta?.content) { + parts.push({ text: delta.content }); + } + if (Array.isArray(delta?.tool_calls)) { + for (const tc of delta.tool_calls) { + const idx = tc.index ?? 0; + const existing = toolCalls.get(idx) ?? { arguments: '' }; + if (tc.id) existing.id = tc.id; + if (tc.function?.name) existing.name = tc.function.name; + if (tc.function?.arguments) { + existing.arguments += tc.function.arguments; + } + toolCalls.set(idx, existing); + } + } + + const finishReason = mapFinishReason(choice.finish_reason); + if (finishReason) { + // Flush any accumulated tool calls together with the finish + // chunk so downstream consumers see calls before/with Finished. + parts.push(...flushToolCallParts()); + sentFinish = true; + yield self.makeResponse( + parts, + finishReason, + usage, + responseId, + modelVersion, + ); + return; + } + + if (parts.length > 0) { + yield self.makeResponse( + parts, + undefined, + undefined, + responseId, + modelVersion, + ); + } + }; + + try { + streamLoop: while (true) { + const { done, value } = await reader.read(); + if (done) break; + buffer += decoder.decode(value, { stream: true }); + const lines = buffer.split('\n'); + buffer = lines.pop() ?? ''; + + for (const line of lines) { + const trimmed = line.trim(); + if (!trimmed.startsWith('data:')) continue; + const data = trimmed.slice(5).trim(); + if (data === '[DONE]') { + break streamLoop; + } + let json: Record; + try { + json = JSON.parse(data) as Record; + } catch { + continue; // Ignore malformed keep-alive/partial lines. + } + yield* handleData(json); + } + } + } finally { + reader.releaseLock(); + try { + await res.body?.cancel(); + } catch { + // Stream already closed. + } + } + + // Stream ended without an explicit finish chunk, or tool calls are + // still pending: emit a final chunk so geminiChat always sees a + // finish reason (it throws InvalidStreamError otherwise). + const trailingParts = flushToolCallParts(); + if (!sentFinish || trailingParts.length > 0) { + yield self.makeResponse( + trailingParts, + FinishReason.STOP, + usage, + responseId, + modelVersion, + ); + } else if (usage) { + // Usage arrived after the finish chunk (stream_options.include_usage): + // surface it in a trailing metadata-only chunk. + const response = new GenerateContentResponse(); + response.candidates = []; + response.usageMetadata = { + promptTokenCount: usage.prompt_tokens ?? 0, + candidatesTokenCount: usage.completion_tokens ?? 0, + totalTokenCount: usage.total_tokens ?? 0, + }; + if (responseId) response.responseId = responseId; + if (modelVersion) response.modelVersion = modelVersion; + yield response; + } + } + + return makeStream(); + } + + /** + * SGLang has no public tokenizer endpoint wired up here; approximate + * (~4 chars/token) so context-compression heuristics still function. + */ + async countTokens( + request: CountTokensParameters, + ): Promise { + let chars = 0; + try { + chars = JSON.stringify(request.contents ?? '').length; + } catch { + chars = 0; + } + return { totalTokens: Math.ceil(chars / 4) }; + } + + async embedContent( + _request: EmbedContentParameters, + ): Promise { + throw new Error( + 'Embeddings are not supported by the SGLang content generator.', + ); + } +} diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 541c16a8f3f..83f4df6c63d 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -36,6 +36,7 @@ export * from './commands/types.js'; export * from './core/baseLlmClient.js'; export * from './core/client.js'; export * from './core/contentGenerator.js'; +export * from './core/sglangContentGenerator.js'; export * from './core/fakeContentGenerator.js'; export * from './core/loggingContentGenerator.js'; export * from './core/geminiChat.js';