diff --git a/AGENTS.md b/AGENTS.md index 93f4bd6..7698995 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -108,6 +108,10 @@ Implemented: - `ti fs delete-file-system` - `ti fs list-file-systems` - `ti fs describe-file-system` +- `ti fs describe-file-system-extract-configuration` +- `ti fs update-file-system-extract-configuration` +- `ti fs describe-file-system-embedding-configuration` +- `ti fs update-file-system-embedding-configuration` - `ti fs check-file-system` - `ti fs copy-file` - `ti fs read-file` @@ -160,6 +164,9 @@ Implemented: - direct TiDB Cloud Filesystem tenant create/list/describe/delete control plane, authoritative display metadata, quota output, and inventory filters from `docs/spec/done/0031-fs-tenant-metadata-control-plane.md` +- optional tenant-scoped media extraction and embedding provider configuration, + provider-secret redaction, and media quota output from + `docs/spec/done/0032-file-system-ai-provider-configuration.md` - ti fs/fs-git/fs-journal/fs-vault commands routed through the bundled `ti-drive9` companion, with ti-owned profile loading, credential storage, region resolution, and output/error handling @@ -341,6 +348,7 @@ internal/db/sqlsingle/ one-statement validation internal/db/validate/ DB flag and request validation helpers internal/dryrun/ shared dry-run result envelope internal/fs/ ti fs control-plane, data-plane, and mount use cases +internal/fs/aiconfig/ optional Filesystem extraction and embedding configuration internal/fs/fscred/ ID-keyed ti fs credentials, selection, and legacy migration internal/fs/mountlocator/ non-secret Drive9 background mount routing state internal/fs/tokenmgmt/ Filesystem token lifecycle and local rotation safety @@ -421,6 +429,11 @@ Follow these rules unless `docs/priciples.md` is updated: - Filesystem token refresh is bearer-only and non-idempotent. Do not retry it after an ambiguous network failure. Generate, list, enable, disable, and delete use only TiDB Cloud public/private keys. +- Filesystem AI provider configuration is optional and uses only TiDB Cloud + public/private keys. Provider secrets come only from + `TI_FS_AI_PROVIDER_API_KEY`, are never persisted locally, and must be + redacted from every output and diagnostic path. Provider-validating PUTs are + not retried after ambiguous failures. - Reject refresh, disable, or deletion of a token correlated with a known active local mount. The error must show exact drain and unmount commands. - Read-only commands reject `--dry-run`. @@ -552,6 +565,10 @@ Implemented command behavior: - `ti fs delete-file-system --file-system-id --dry-run` - `ti fs list-file-systems` - `ti fs describe-file-system --file-system-id ` +- `ti fs describe-file-system-extract-configuration --file-system-id --media-type image` +- `ti fs update-file-system-extract-configuration --file-system-id --media-type image --enabled false` +- `ti fs describe-file-system-embedding-configuration --file-system-id ` +- `ti fs update-file-system-embedding-configuration --file-system-id --enabled false` - `ti fs check-file-system` - `ti fs check-file-system --file-system-id ` - `ti fs copy-file --from-local ./README.md --to-remote /workspace/README.md` diff --git a/README.md b/README.md index 07079fa..1bd01c2 100644 --- a/README.md +++ b/README.md @@ -222,6 +222,24 @@ ti fs generate-file-system-scoped-token \ An owner FS token authorizes Filesystem use and token management, but it is not a TiDB Cloud administrative credential. Creating, listing, describing, and deleting Filesystem resources require TiDB Cloud API keys. In particular, `ti fs delete-file-system` always requires an explicit `--file-system-id`; `TI_FS_TOKEN` cannot select or authorize deletion of the Filesystem itself. For token list, enable, disable, and delete commands, `--file-system-id` is required only when the command uses TiDB Cloud API keys. When an owner token is supplied through `--fs-token` or `TI_FS_TOKEN`, `ti` derives the Filesystem ID from that token. +Filesystem AI provider configuration is optional. Without it, existing create, mount, file, search, layer, vault, journal, and Git workflows continue to work normally. Organization administrators can inspect or customize image, audio, and video extraction and OpenAI-compatible 1024-dimension embedding for an explicit Filesystem ID: + +```shell +ti fs describe-file-system-extract-configuration \ + --file-system-id "$FILE_SYSTEM_ID" \ + --media-type image + +TI_FS_AI_PROVIDER_API_KEY="" \ +ti fs update-file-system-extract-configuration \ + --file-system-id "$FILE_SYSTEM_ID" \ + --media-type image \ + --enabled true \ + --provider-api-base https://api.openai.com/v1 \ + --provider-model "" +``` + +Extraction supports the `openai` protocol for image, audio, and video. Audio also supports Alibaba Cloud Model Studio Qwen ASR through `--provider-protocol qwen-asr`. Exact OpenAI-compatible endpoints can work when they implement the required API contract; native Anthropic, Gemini, Vertex AI, Bedrock, and Azure OpenAI interfaces are not supported. Enabling or replacing a provider sends a real validation request that can incur a small provider charge. After enablement, Filesystem content is sent to the selected extraction provider, and text or descriptions are sent to the selected embedding provider. The provider key is accepted only through `TI_FS_AI_PROVIDER_API_KEY`, is not stored locally, and is returned by the backend only in masked form. + Generation does not modify local credentials by default. Add `--store-locally` to select the new token locally; if a selected token already exists, add `--replace` explicitly. Replacing local selection does not revoke the previous remote token. Use immutable `token_id` values from the list response to disable, enable, or permanently revoke a token: ```shell @@ -295,6 +313,10 @@ ti fs refresh-file-system-token ti fs delete-file-system ti fs list-file-systems ti fs describe-file-system +ti fs describe-file-system-extract-configuration +ti fs update-file-system-extract-configuration +ti fs describe-file-system-embedding-configuration +ti fs update-file-system-embedding-configuration ti fs check-file-system ti fs copy-file ti fs read-file diff --git a/docs/pingcap-docs/docs b/docs/pingcap-docs/docs index 0a65ec5..07c3776 160000 --- a/docs/pingcap-docs/docs +++ b/docs/pingcap-docs/docs @@ -1 +1 @@ -Subproject commit 0a65ec54f3f150acb7ca4c39dee1a2d1f0394e4e +Subproject commit 07c3776792ff8b135608e49ed7bb4100aee83e38 diff --git a/docs/spec/0032-homebrew-and-scoop-distribution.md b/docs/spec/0033-homebrew-and-scoop-distribution.md similarity index 100% rename from docs/spec/0032-homebrew-and-scoop-distribution.md rename to docs/spec/0033-homebrew-and-scoop-distribution.md diff --git a/docs/spec/0033-serverless-function-deployment.md b/docs/spec/0034-serverless-function-deployment.md similarity index 100% rename from docs/spec/0033-serverless-function-deployment.md rename to docs/spec/0034-serverless-function-deployment.md diff --git a/docs/spec/done/0012-install-and-update-distribution.md b/docs/spec/done/0012-install-and-update-distribution.md index c1ea241..8a602c7 100644 --- a/docs/spec/done/0012-install-and-update-distribution.md +++ b/docs/spec/done/0012-install-and-update-distribution.md @@ -2,7 +2,7 @@ ## Goal -Make `tdc` installable and updatable through deterministic GitHub Releases artifacts. The MVP channel is GoReleaser plus GitHub Releases, with shell and PowerShell installers. Homebrew and Scoop are intentionally deferred to `docs/spec/0032-homebrew-and-scoop-distribution.md`. +Make `tdc` installable and updatable through deterministic GitHub Releases artifacts. The MVP channel is GoReleaser plus GitHub Releases, with shell and PowerShell installers. Homebrew and Scoop are intentionally deferred to `docs/spec/0033-homebrew-and-scoop-distribution.md`. ## User-facing Commands @@ -266,7 +266,7 @@ Installer scripts: - Silent auto-update. - Updating TiDB Cloud credentials or DB SQL credentials. - Config migrations that modify user config during update. -- Homebrew tap and Scoop bucket publishing. See `0032-homebrew-and-scoop-distribution.md`. +- Homebrew tap and Scoop bucket publishing. See `0033-homebrew-and-scoop-distribution.md`. - Linux apt/yum repositories. - Winget publishing. - Notarization or binary signing beyond SHA-256 checksums for MVP. diff --git a/docs/spec/done/0032-file-system-ai-provider-configuration.md b/docs/spec/done/0032-file-system-ai-provider-configuration.md new file mode 100644 index 0000000..d2ae322 --- /dev/null +++ b/docs/spec/done/0032-file-system-ai-provider-configuration.md @@ -0,0 +1,444 @@ +# File System AI Provider Configuration + +## Goal + +Expose the TiDB Cloud Filesystem server's tenant-scoped media extraction and embedding provider configuration through `ti`. A user with TiDB Cloud organization credentials can inspect, enable, update, and disable the AI processing configuration of one File System without invoking the Drive9 admin CLI directly. + +This spec also extends existing File System inventory output with the media and video quota fields already returned by the server. It is a `ti` client integration of the backend behavior introduced by `tidbcloud/fs` and Drive9 changes corresponding to media extraction configuration, media quota accounting, and embedding configuration. It does not add AI execution to `ti` itself. + +## Scope + +Add four commands: + +```text +ti fs describe-file-system-extract-configuration +ti fs update-file-system-extract-configuration +ti fs describe-file-system-embedding-configuration +ti fs update-file-system-embedding-configuration +``` + +Extend the quota objects returned by these existing commands: + +```text +ti fs list-file-systems +ti fs describe-file-system +``` + +The new commands call the hosted File System control-plane APIs directly through typed clients under `internal/api/fs`. They must not shell out to `ti-drive9 admin`, parse companion output, or make `ref/` a runtime, build, or test dependency. + +## User-Facing Commands + +Inspect image extraction configuration: + +```bash +ti fs describe-file-system-extract-configuration \ + --file-system-id \ + --media-type image +``` + +Enable image extraction with OpenAI: + +```bash +TI_FS_AI_PROVIDER_API_KEY="" \ +ti fs update-file-system-extract-configuration \ + --file-system-id \ + --media-type image \ + --enabled true \ + --provider-api-base https://api.openai.com/v1 \ + --provider-model \ + --provider-protocol openai +``` + +Enable OpenAI-compatible audio transcription: + +```bash +TI_FS_AI_PROVIDER_API_KEY="" \ +ti fs update-file-system-extract-configuration \ + --file-system-id \ + --media-type audio \ + --enabled true \ + --provider-api-base https://api.openai.com/v1 \ + --provider-model gpt-4o-transcribe \ + --provider-protocol openai +``` + +Enable Alibaba Cloud Model Studio Qwen ASR: + +```bash +TI_FS_AI_PROVIDER_API_KEY="" \ +ti fs update-file-system-extract-configuration \ + --file-system-id \ + --media-type audio \ + --enabled true \ + --provider-api-base https://dashscope.aliyuncs.com/compatible-mode/v1 \ + --provider-model qwen3-asr-flash \ + --provider-protocol qwen-asr +``` + +Update only the extraction prompt of an already enabled configuration: + +```bash +ti fs update-file-system-extract-configuration \ + --file-system-id \ + --media-type image \ + --prompt "Describe the image and return searchable attributes." +``` + +Disable extraction for one media type: + +```bash +ti fs update-file-system-extract-configuration \ + --file-system-id \ + --media-type image \ + --enabled false +``` + +Inspect embedding configuration: + +```bash +ti fs describe-file-system-embedding-configuration \ + --file-system-id +``` + +Enable app-managed embedding: + +```bash +TI_FS_AI_PROVIDER_API_KEY="" \ +ti fs update-file-system-embedding-configuration \ + --file-system-id \ + --enabled true \ + --provider-api-base https://api.openai.com/v1 \ + --provider-model text-embedding-3-small +``` + +Disable app-managed embedding: + +```bash +ti fs update-file-system-embedding-configuration \ + --file-system-id \ + --enabled false +``` + +All four commands require an explicit `--file-system-id`. They do not infer the target from an FS token or a historical profile default. The global `--region` precedence remains unchanged. + +## Provider Support + +The backend selects an implementation by wire protocol, not by a vendor enum. `ti` therefore exposes `--provider-protocol` only where the backend API has a protocol field and does not add `--provider openai|alibaba`. + +| Provider or interface | Embedding | Image extraction | Audio extraction | Video extraction | +| --- | --- | --- | --- | --- | +| OpenAI | Supported through `openai` | Supported through `openai` | Supported through `openai` | Supported through `openai` | +| OpenAI-compatible API | Supported when the endpoint satisfies the exact contract and returns 1024 dimensions | Supported when the endpoint implements the required Chat Completions vision contract | Supported when the endpoint implements the required transcription contract | Supported when the endpoint implements the required Chat Completions vision contract | +| Alibaba Cloud Model Studio Qwen ASR | No native provider contract | No native provider contract | Supported through `qwen-asr` | No native provider contract | + +Supported protocols and request shapes are: + +- `openai` embedding: `POST /v1/embeddings` with a fixed requested and returned vector width of 1024. +- `openai` image and video extraction: `POST /v1/chat/completions` with image input. +- `openai` audio extraction: `POST /v1/audio/transcriptions`. +- `qwen-asr` audio extraction: Alibaba Cloud Model Studio compatible-mode `POST /v1/chat/completions` with Qwen ASR audio input. + +For extraction, `--provider-protocol` defaults to `openai`. Image and video accept only `openai`; audio accepts `openai` and `qwen-asr`. Embedding is always OpenAI-compatible in this version and therefore does not expose `--provider-protocol`. + +Anthropic Messages, Google Gemini or Vertex AI native APIs, AWS Bedrock native APIs, and Azure OpenAI native authentication and deployment URL forms are not supported. A provider not listed above may work only if it exposes the exact OpenAI-compatible request and response contract and passes server-side validation. Documentation must describe that as conditional compatibility, not as officially supported vendor integration. + +## Authentication And Secret Input + +The four configuration commands use TiDB Cloud public/private API credentials and organization-level authorization. An FS owner or scoped token is neither required nor sufficient for these admin operations. + +The provider API key is accepted only through: + +```text +TI_FS_AI_PROVIDER_API_KEY +``` + +Do not add a plaintext `--provider-api-key` flag because command-line arguments are commonly retained in shell history and process listings. Do not persist the provider key in `~/.ti/config`, `~/.ti/credentials`, `~/.ti/.preferences`, mount locators, or another local file. + +The environment variable is read only when an update supplies a provider credential set. It is ignored for disable and prompt-only operations, so an ambient value cannot accidentally turn a disable request into a provider replacement. + +The provider key is sent in the HTTPS request body to the File System backend. The backend validates it against the provider and stores it encrypted. GET responses return only a masked key. `ti` must register the plaintext with the existing redactor before request encoding so it cannot appear in debug output, operation logs, telemetry, dry-run output, errors, or test failure diagnostics. + +Provider endpoints receive File System content for extraction and text or descriptions for embedding after the configuration is enabled. User documentation must state this data-sharing boundary and tell users to choose a provider account and retention policy appropriate for their data. + +## Extract Configuration Contract + +Supported media types are exactly: + +```text +image +audio +video +``` + +`text` is reserved by the backend storage model but is not an available API media type. `ti` rejects `--media-type text` and unknown values locally. + +The read command calls: + +```http +GET /v1/admin/tenants/{file_system_id}/extract-config/{media_type} +``` + +The update command calls: + +```http +PUT /v1/admin/tenants/{file_system_id}/extract-config/{media_type} +Content-Type: application/json +``` + +The typed response contains: + +```json +{ + "enabled": true, + "api_base": "https://api.openai.com/v1", + "api_key": "sk-a********", + "model": "", + "protocol": "openai", + "prompt": "", + "source": "custom", + "updated_at": "2026-08-26T10:00:00Z" +} +``` + +`source` is server-owned and can be `custom`, `default`, or `none`. `ti` preserves unknown future source values rather than rejecting an otherwise valid response. + +Extract updates preserve the backend's partial-update semantics: + +- `--enabled true` on an absent or disabled custom configuration requires `--provider-api-base`, `--provider-model`, and `TI_FS_AI_PROVIDER_API_KEY`. +- Replacing any member of that provider trio requires all three members in the same command. +- `--enabled false` sends only `enabled=false` and clears the custom provider fields on the server. +- `--prompt` can update the prompt of an already enabled configuration without resending credentials. +- An explicit empty `--prompt ""` clears the custom prompt and returns to the backend's default prompt behavior. +- Changing `--provider-protocol` for audio requires the complete provider trio. +- A command that supplies none of `--enabled`, `--provider-api-base`, `--provider-model`, `--provider-protocol`, or `--prompt` is a usage error. +- Explicit empty provider base, model, or key values are invalid. + +`--enabled` requires an explicit `true` or `false` value. It must retain an unset state internally so omission is distinguishable from `false`. + +## Embedding Configuration Contract + +The read command calls: + +```http +GET /v1/admin/tenants/{file_system_id}/embedding-config +``` + +The update command calls: + +```http +PUT /v1/admin/tenants/{file_system_id}/embedding-config +Content-Type: application/json +``` + +The typed response contains: + +```json +{ + "enabled": true, + "api_base": "https://api.openai.com/v1", + "api_key": "sk-a********", + "model": "text-embedding-3-small", + "source": "custom", + "generation": 3, + "updated_at": "2026-08-27T10:00:00Z" +} +``` + +`source` can be `custom`, `default`, `none`, or `database_auto`. `generation` is server-owned and changes when the vector contract changes. `ti` must not synthesize or mutate either field. + +Embedding PUT is a full replacement: + +- Enabling requires `--enabled true`, `--provider-api-base`, `--provider-model`, and `TI_FS_AI_PROVIDER_API_KEY` in the same command. +- Disabling requires `--enabled false` and rejects provider configuration flags. +- The vector width is fixed at 1024. Do not expose a dimensions flag. +- Native File Systems using TiDB database-managed `auto` embedding return `source=database_auto`; attempting app-managed configuration is not applicable and the backend returns HTTP 409. +- Shared File Systems and native File Systems whose effective embedding mode is `fts_only` support app-managed configuration. + +## Provider Validation And Retry Boundary + +When provider credentials are created, re-enabled, or replaced, the backend performs a real provider request before persisting the new configuration: + +- Embedding validation requires exactly 1024 finite vector values. +- Image and video validation send a small built-in image through the configured vision model. +- Audio validation sends a bounded silent audio sample through the selected protocol. + +The validation can incur a small charge on the user's provider account. `ti` must mention this in command help and documentation. + +The mutating commands must not automatically retry a PUT after an ambiguous network failure, timeout, or lost response. A provider request may already have been charged and the configuration may already have been persisted. The error should direct the user to run the matching describe command before deciding whether to retry. + +`ti` validates required fields, media/protocol compatibility, HTTPS URL syntax, and prompt size before sending the request. The server remains authoritative for endpoint safety, DNS resolution, provider capability, credentials, model compatibility, and response validation. `ti` never calls the provider directly. + +## Quota And Usage Output + +Extend `internal/api/fs.AdminTenantQuotaConfig` with: + +```json +{ + "max_media_llm_files": 100, + "max_video_llm_files": 30 +} +``` + +Extend `internal/api/fs.AdminTenantQuotaUsage` with: + +```json +{ + "media_file_count": 12, + "video_file_count": 3 +} +``` + +The exact values are server-owned and may vary by plan or deployment. `ti` must not hardcode the example limits as product defaults. The new fields appear in JSON output from `list-file-systems` and `describe-file-system` whenever the backend returns quota data. + +The text list remains compact and does not add four quota columns. Text describe renders the new values with stable labels alongside the existing quota details. This spec does not expose an admin command for changing quota values. + +## Output + +Configuration commands return structured JSON by default and support `--output text` and `--query` through the shared control-plane rendering path. Output field names match the backend contract except that the public resource selector remains `file_system_id` rather than `tenant_id`. + +Example extract text output: + +```text +File system ID: +Media type: image +Enabled: true +Source: custom +Provider API base: https://api.openai.com/v1 +Provider API key: sk-a******** +Provider model: +Provider protocol: openai +Prompt: +Updated at: 2026-08-26T10:00:00Z +``` + +Example embedding text output: + +```text +File system ID: +Enabled: true +Source: custom +Provider API base: https://api.openai.com/v1 +Provider API key: sk-a******** +Provider model: text-embedding-3-small +Generation: 3 +Updated at: 2026-08-27T10:00:00Z +``` + +Missing optional fields render as `none`, not an empty unlabeled value. JSON preserves omission from the backend where practical. Neither output mode ever returns the plaintext provider key. + +## Dry Run + +Both update commands support `--dry-run`. Dry run performs local profile, TiDB Cloud credential, region, File System ID, media type, protocol, URL, required-field, and prompt validation. It does not call the File System backend or the external provider and does not write local state. + +Dry-run output includes the target API path, non-secret request fields, and whether a provider API key was supplied. It renders the key only as a fixed redacted marker and never includes a masked prefix derived from the plaintext. + +Describe commands are read-only and reject `--dry-run`. + +## Errors + +Use stable typed errors and preserve a safe backend reason: + +| Condition | Behavior | +| --- | --- | +| Missing or invalid File System ID | Usage error before API execution | +| Unsupported media type or protocol | Usage error listing accepted values | +| Missing provider trio while enabling or replacing | Usage error naming all required inputs | +| Non-HTTPS provider API base | Usage error before API execution | +| Missing TiDB Cloud credentials | Existing authentication-required error and API-key creation URL | +| HTTP 401 | Existing TiDB Cloud authentication error | +| HTTP 403 | Existing organization authorization error | +| HTTP 404 | Existing File System resource-not-found behavior | +| HTTP 409 for database-managed embedding | `fs.embedding_configuration_not_applicable` with `source=database_auto` guidance | +| Provider rejects credentials, endpoint, model, or media capability | `fs.ai_provider_validation_failed` with the backend's sanitized reason | +| Provider rate limit, timeout, or availability failure | Runtime error that warns the user not to retry blindly and points to the describe command | +| Malformed successful response | API response-contract error without exposing credentials or provider payloads | + +Do not include provider response bodies, headers, file content, prompts, or API keys in errors. The server's bounded validation category and HTTP status may be retained. + +## Implementation Design + +- Add typed extract and embedding request/response models and methods under `internal/api/fs`. +- Add a focused service package under `internal/fs` for AI configuration validation and use cases; do not place provider HTTP implementations in `ti`. +- Register all four commands through `controlPlaneCommandSpec` so auth, dry-run, query, JSON/text rendering, and error behavior remain consistent. +- Add explicit static permissions for extract configuration read/update and embedding configuration read/update. Do not infer permission from command names. +- Reuse the existing canonical File System endpoint resolver and Digest-authenticated TiDB Cloud profile loading. +- Add the provider API key to the API client's redactor before any request can be logged. +- Do not invoke or depend on the bundled companion for these four commands. +- Do not add local configuration files, credential fields, context entries, caches, or migration work. +- Update `README.md`, generated command documentation, PingCAP Preview documentation, `AGENTS.md`, and command-surface tests when implementation begins. + +## Logging And Telemetry + +Operation logs and telemetry may record: + +- Command path. +- Region code. +- Media type as a bounded enum. +- Provider protocol as a bounded enum. +- Enabled/disabled action. +- Stable error code, exit code, and duration. + +They must not record provider API base, API key, model, prompt, File System ID, provider response, extracted content, embedding input, or returned vectors. Flag names may be recorded under the existing policy, but flag values remain excluded. + +## Test Plan + +Unit tests must cover: + +- Typed request paths, methods, headers, and JSON bodies for all four APIs. +- Extract partial updates and embedding full replacements. +- Explicit `--enabled true|false` handling and omission detection. +- Media type and protocol compatibility, including local rejection of `text`. +- Required provider trio and environment-only secret loading. +- HTTPS provider URL validation and the 8 KiB prompt limit. +- Redaction from errors, debug output, operation logs, dry-run output, and test diagnostics. +- Response decoding for every documented `source` value, masked keys, generation, and missing optional fields. +- HTTP 409 database-managed embedding behavior and provider-validation error mappings. +- No automatic retry of update requests. +- JSON, text, and `--query` rendering. +- The four new quota and usage fields in list and describe responses. + +Black-box e2e uses a fake File System API and must verify complete command wiring without importing fixtures or code from `ref/`. + +Live e2e must at minimum: + +1. Create a uniquely named temporary File System through the existing lifecycle. +2. Describe all three extract media configurations and the embedding configuration. +3. Disable an applicable configuration, describe it again, and verify the server response without requiring an external provider account. +4. Verify the new quota and usage fields are preserved when returned. +5. Delete only the temporary File System through the existing cleanup path. + +Successful enablement against a paid external provider is separately opt-in because it requires a provider credential and can incur charges. When configured, it must use dedicated ignored environment variables and must not become a prerequisite for `make test`, `make e2e`, or the ordinary live-e2e suite. + +Before release, probe the read endpoints in all four supported File System regions. A region that has not deployed the required backend API must fail as unavailable; `ti` must not fall back to another region or emulate the operation through Drive9 local state. + +## Acceptance Criteria + +- Users can inspect and update image, audio, and video extraction configuration for one explicit File System. +- Users can inspect and update app-managed embedding configuration where the server reports it is applicable. +- Help and documentation clearly distinguish OpenAI, conditional OpenAI-compatible APIs, and Alibaba Cloud Model Studio Qwen ASR support. +- Unsupported native provider interfaces are not presented as supported. +- Provider API keys are environment-only, never stored locally, always redacted, and returned only in masked server form. +- Provider validation and its possible charge are disclosed, and mutating requests are not automatically retried after ambiguous failures. +- Read and update commands use direct typed File System APIs with TiDB Cloud credentials and never invoke `ti-drive9 admin`. +- Existing FS token-only data-plane workflows remain unchanged and cannot perform these organization-level operations. +- `list-file-systems` and `describe-file-system` preserve media/video quota limits and usage counters returned by the backend. +- Unit, black-box e2e, and applicable live-e2e coverage pass without depending on `ref/`. + +## Out Of Scope + +- Implementing embedding, image understanding, transcription, video processing, vector search, or provider validation inside `ti`. +- Supporting text extraction configuration. +- Supporting native Anthropic, Gemini, Vertex AI, Bedrock, or Azure OpenAI interfaces. +- Configurable embedding dimensions or separate providers for document and query embedding. +- Changing File System quota values from `ti`. +- Listing provider models or auto-detecting a provider from its model name or URL. +- Persisting provider credentials locally or adding an interactive credential prompt. +- Backfilling or reprocessing historical files from `ti`. +- Falling back across regions, providers, protocols, or File Systems. + +## Dependencies + +- `docs/spec/done/0031-fs-tenant-metadata-control-plane.md` +- Hosted File System extract-config endpoints for `image`, `audio`, and `video`. +- Hosted File System embedding-config endpoint. +- Hosted File System inventory responses containing media/video quota limits and usage counters. diff --git a/e2e/cli_test.go b/e2e/cli_test.go index d359194..daf00df 100644 --- a/e2e/cli_test.go +++ b/e2e/cli_test.go @@ -98,6 +98,16 @@ func TestHelpAndVersion(t *testing.T) { listFileSystems.wantExitCode(0) listFileSystems.wantStdoutContains("[--display-name ]") listFileSystems.wantStdoutContains("[--label ]") + updateEmbedding := runTI(t, bin, "fs", "update-file-system-embedding-configuration", "help") + updateEmbedding.wantExitCode(0) + updateEmbedding.wantStdoutContains("--enabled (required)") + updateEmbedding.wantStdoutContains("--file-system-id (required)") + updateEmbedding.wantStdoutNotContains("--fs-token") + updateExtract := runTI(t, bin, "fs", "update-file-system-extract-configuration", "help") + updateExtract.wantExitCode(0) + updateExtract.wantStdoutContains("[--enabled ]") + updateExtract.wantStdoutContains("--media-type (required)") + updateExtract.wantStdoutNotContains("--provider-api-key") createDBCluster := runTI(t, bin, "db", "create-db-cluster", "help") createDBCluster.wantExitCode(0) @@ -786,6 +796,10 @@ func TestFSRemoteInventoryAndIDCredentialSelectionAcrossCommandFamilies(t *testi describe.wantStdoutContains(`"display_name": "tenant-aws-us-west-2"`) describe.wantStdoutContains(`"labels": {}`) describe.wantStdoutContains(`"quota": {`) + describe.wantStdoutContains(`"max_media_llm_files": 100`) + describe.wantStdoutContains(`"max_video_llm_files": 30`) + describe.wantStdoutContains(`"media_file_count": 0`) + describe.wantStdoutContains(`"video_file_count": 0`) describe.wantStdoutContains(`"region_code": "aws-us-west-2"`) describe.wantStdoutNotContains("drive9_") textDescribe := runTIWithInput(t, bin, "", baseEnv, "--profile", "stage", "--region", "aws-us-west-2", "fs", "describe-file-system", "--file-system-id", "tenant-aws-us-west-2", "--output", "text") @@ -793,7 +807,70 @@ func TestFSRemoteInventoryAndIDCredentialSelectionAcrossCommandFamilies(t *testi textDescribe.wantStdoutContains("File system ID: tenant-aws-us-west-2") textDescribe.wantStdoutContains("Display name: tenant-aws-us-west-2") textDescribe.wantStdoutContains("Labels: none") + textDescribe.wantStdoutContains("Quota max media LLM files: 100") + textDescribe.wantStdoutContains("Usage video file count: 0") textDescribe.wantStdoutNotContains(`"file_system_id"`) + + describeExtract := runTIWithInput(t, bin, "", baseEnv, "--profile", "stage", "--region", "aws-us-west-2", "fs", "describe-file-system-extract-configuration", "--file-system-id", "tenant-aws-us-west-2", "--media-type", "image") + describeExtract.wantExitCode(0) + describeExtract.wantStdoutContains(`"file_system_id": "tenant-aws-us-west-2"`) + describeExtract.wantStdoutContains(`"media_type": "image"`) + describeExtract.wantStdoutContains(`"source": "none"`) + describeExtractText := runTIWithInput(t, bin, "", baseEnv, "--profile", "stage", "--region", "aws-us-west-2", "fs", "describe-file-system-extract-configuration", "--file-system-id", "tenant-aws-us-west-2", "--media-type", "audio", "--output", "text") + describeExtractText.wantExitCode(0) + describeExtractText.wantStdoutContains("Provider API key: none") + describeExtractQuery := runTIWithInput(t, bin, "", baseEnv, "--profile", "stage", "--region", "aws-us-west-2", "fs", "describe-file-system-extract-configuration", "--file-system-id", "tenant-aws-us-west-2", "--media-type", "video", "--query", "source", "--output", "text") + describeExtractQuery.wantExitCode(0) + if strings.TrimSpace(describeExtractQuery.stdout) != "none" { + describeExtractQuery.fail("extract query output mismatch") + } + describeEmbedding := runTIWithInput(t, bin, "", baseEnv, "--profile", "stage", "--region", "aws-us-west-2", "fs", "describe-file-system-embedding-configuration", "--file-system-id", "tenant-aws-us-west-2") + describeEmbedding.wantExitCode(0) + describeEmbedding.wantStdoutContains(`"source": "none"`) + + requestsBeforeAIDryRun := westControl.requestCount() + aiEnv := append(baseEnv, "TI_FS_AI_PROVIDER_API_KEY=provider-e2e-secret") + dryRunExtract := runTIWithInput(t, bin, "", aiEnv, "--profile", "stage", "--region", "aws-us-west-2", "fs", "update-file-system-extract-configuration", + "--file-system-id", "tenant-aws-us-west-2", "--media-type", "audio", "--enabled", "true", + "--provider-api-base", "https://provider.example/v1", "--provider-model", "audio-model", "--provider-protocol", "qwen-asr", "--dry-run") + dryRunExtract.wantExitCode(0) + dryRunExtract.wantStdoutContains(`"api_key": "[REDACTED]"`) + dryRunExtract.wantStdoutContains(`"provider_api_key_supplied": true`) + dryRunExtract.wantStdoutNotContains("provider-e2e-secret") + if got := westControl.requestCount(); got != requestsBeforeAIDryRun { + t.Fatalf("AI dry-run sent remote requests: before=%d after=%d", requestsBeforeAIDryRun, got) + } + updateExtract := runTIWithInput(t, bin, "", aiEnv, "--profile", "stage", "--region", "aws-us-west-2", "fs", "update-file-system-extract-configuration", + "--file-system-id", "tenant-aws-us-west-2", "--media-type", "image", "--enabled", "true", + "--provider-api-base", "https://provider.example/v1", "--provider-model", "vision-model", "--debug") + updateExtract.wantExitCode(0) + updateExtract.wantStdoutContains(`"api_key": "pro********"`) + updateExtract.wantStdoutNotContains("provider-e2e-secret") + updateExtract.wantStderrNotContains("provider-e2e-secret") + if err := filepath.Walk(filepath.Join(home, ".ti"), func(path string, info os.FileInfo, walkErr error) error { + if walkErr != nil || info.IsDir() { + return walkErr + } + data, readErr := os.ReadFile(path) + if readErr != nil { + return readErr + } + if strings.Contains(string(data), "provider-e2e-secret") { + return fmt.Errorf("provider key persisted in %s", path) + } + return nil + }); err != nil { + t.Fatal(err) + } + updateEmbedding := runTIWithInput(t, bin, "", baseEnv, "--profile", "stage", "--region", "aws-us-west-2", "fs", "update-file-system-embedding-configuration", + "--file-system-id", "tenant-aws-us-west-2", "--enabled", "false", "--query", "enabled") + updateEmbedding.wantExitCode(0) + updateEmbedding.wantStdoutContains("false") + invalidMedia := runTIWithInput(t, bin, "", baseEnv, "--profile", "stage", "fs", "describe-file-system-extract-configuration", "--file-system-id", "tenant-aws-us-east-1", "--media-type", "text") + invalidMedia.wantExitCode(2) + invalidMedia.wantStderrContains("image, audio, or video") + missingEnabledValue := runTIWithInput(t, bin, "", baseEnv, "--profile", "stage", "fs", "update-file-system-embedding-configuration", "--file-system-id", "tenant-aws-us-east-1", "--enabled") + missingEnabledValue.wantExitCode(2) callsBeforeMissingSelectorCommands := len(readFakeDrive9Calls(t, recordPath)) for _, args := range [][]string{ {"fs", "list-files", "--path", "/"}, @@ -1327,6 +1404,57 @@ func (f *fakeFSTenantControlPlane) serveHTTP(w http.ResponseWriter, r *http.Requ } sort.Slice(items, func(i, j int) bool { return items[i].TenantID < items[j].TenantID }) _ = json.NewEncoder(w).Encode(map[string]any{"tenants": items, "page": 1, "page_size": 100, "next_page": 0}) + case strings.Contains(r.URL.Path, "/extract-config/"): + parts := strings.Split(strings.TrimPrefix(r.URL.Path, "/v1/admin/tenants/"), "/") + if len(parts) != 3 || parts[1] != "extract-config" { + http.NotFound(w, r) + return + } + if _, ok := f.tenants[parts[0]]; !ok { + http.Error(w, `{"error":"tenant not found"}`, http.StatusNotFound) + return + } + switch r.Method { + case http.MethodGet: + _ = json.NewEncoder(w).Encode(map[string]any{"enabled": false, "source": "none"}) + case http.MethodPut: + var body map[string]any + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + if key, ok := body["api_key"].(string); ok { + if key != "provider-e2e-secret" { + http.Error(w, `{"error":"unexpected provider key"}`, http.StatusBadRequest) + return + } + body["api_key"] = "pro********" + } + body["source"] = "custom" + _ = json.NewEncoder(w).Encode(body) + default: + http.Error(w, `{"error":"method not allowed"}`, http.StatusMethodNotAllowed) + } + case strings.HasSuffix(r.URL.Path, "/embedding-config"): + id := strings.TrimSuffix(strings.TrimPrefix(r.URL.Path, "/v1/admin/tenants/"), "/embedding-config") + if _, ok := f.tenants[id]; !ok { + http.Error(w, `{"error":"tenant not found"}`, http.StatusNotFound) + return + } + switch r.Method { + case http.MethodGet: + _ = json.NewEncoder(w).Encode(map[string]any{"enabled": false, "source": "none"}) + case http.MethodPut: + var body map[string]any + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + body["source"] = "custom" + _ = json.NewEncoder(w).Encode(body) + default: + http.Error(w, `{"error":"method not allowed"}`, http.StatusMethodNotAllowed) + } case strings.HasPrefix(r.URL.Path, "/v1/admin/tenants/"): id := strings.TrimPrefix(r.URL.Path, "/v1/admin/tenants/") tenant, ok := f.tenants[id] @@ -1378,8 +1506,8 @@ func (f *fakeFSTenantControlPlane) hasRequest(method, path string, queryParts .. func fakeFSTenantQuota() map[string]any { return map[string]any{ - "config": map[string]any{"max_storage_size": 1024, "max_file_size": 128, "max_file_count": 1000, "tidbcloud_spending_limit": nil}, - "usage": map[string]any{"storage_bytes": 0, "reserved_bytes": 0, "file_count": 0}, + "config": map[string]any{"max_storage_size": 1024, "max_file_size": 128, "max_file_count": 1000, "max_media_llm_files": 100, "max_video_llm_files": 30, "tidbcloud_spending_limit": nil}, + "usage": map[string]any{"storage_bytes": 0, "reserved_bytes": 0, "file_count": 0, "media_file_count": 0, "video_file_count": 0}, } } diff --git a/e2e/live_test.go b/e2e/live_test.go index fe4cc1e..8095914 100644 --- a/e2e/live_test.go +++ b/e2e/live_test.go @@ -156,8 +156,28 @@ func TestLiveFSRemoteInventoryLifecycle(t *testing.T) { describe.wantStdoutContains(`"display_name": "` + displayName + `"`) describe.wantStdoutContains(`"` + labelKey + `": "` + labelValue + `"`) describe.wantStdoutContains(`"quota":`) + describe.wantStdoutContains(`"max_media_llm_files":`) + describe.wantStdoutContains(`"max_video_llm_files":`) + describe.wantStdoutContains(`"media_file_count":`) + describe.wantStdoutContains(`"video_file_count":`) describe.wantStdoutContains(`"has_local_token": true`) describe.wantStdoutNotContains(created.FSToken) + for _, mediaType := range []string{"image", "audio", "video"} { + extract := runTI(t, bin, "--profile", profileName, "fs", "describe-file-system-extract-configuration", "--file-system-id", created.FileSystemID, "--media-type", mediaType) + extract.wantExitCode(0) + extract.wantStdoutContains(`"file_system_id": "` + created.FileSystemID + `"`) + extract.wantStdoutContains(`"media_type": "` + mediaType + `"`) + extract.wantStdoutContains(`"source":`) + } + embedding := runTI(t, bin, "--profile", profileName, "fs", "describe-file-system-embedding-configuration", "--file-system-id", created.FileSystemID) + embedding.wantExitCode(0) + embedding.wantStdoutContains(`"file_system_id": "` + created.FileSystemID + `"`) + embedding.wantStdoutContains(`"source":`) + disableExtract := runTI(t, bin, "--profile", profileName, "fs", "update-file-system-extract-configuration", "--file-system-id", created.FileSystemID, "--media-type", "image", "--enabled", "false") + disableExtract.wantExitCode(0) + disabledExtract := runTI(t, bin, "--profile", profileName, "fs", "describe-file-system-extract-configuration", "--file-system-id", created.FileSystemID, "--media-type", "image") + disabledExtract.wantExitCode(0) + disabledExtract.wantStdoutContains(`"enabled": false`) } func TestLiveCLICommandSurface(t *testing.T) { @@ -240,6 +260,8 @@ func TestLiveFSCommandSurface(t *testing.T) { {"fs", "help"}, {"fs", "create-file-system", "help"}, {"fs", "import-file-system-token", "help"}, {"fs", "list-file-systems", "help"}, {"fs", "describe-file-system", "help"}, {"fs", "copy-file", "help"}, + {"fs", "describe-file-system-extract-configuration", "help"}, {"fs", "update-file-system-extract-configuration", "help"}, + {"fs", "describe-file-system-embedding-configuration", "help"}, {"fs", "update-file-system-embedding-configuration", "help"}, {"fs", "read-file", "help"}, {"fs", "chmod-file", "help"}, {"fs", "create-symlink", "help"}, {"fs", "create-hardlink", "help"}, {"fs", "create-layer", "help"}, {"fs", "list-layers", "help"}, @@ -260,6 +282,8 @@ func TestLiveFSCommandSurface(t *testing.T) { testLiveMutatingDryRuns(t, bin, profileName, [][]string{ {"fs", "create-file-system", "--display-name", "ti-e2e-dry-run", "--label", "environment=test", "--wait"}, {"fs", "delete-file-system", "--file-system-id", selected.FSTenantID}, + {"fs", "update-file-system-extract-configuration", "--file-system-id", selected.FSTenantID, "--media-type", "image", "--enabled", "false"}, + {"fs", "update-file-system-embedding-configuration", "--file-system-id", selected.FSTenantID, "--enabled", "false"}, {"fs", "create-layer", "--layer-id", "layer-1", "--base-root-path", "/workspace", "--layer-name", "dev"}, {"fs", "create-layer-checkpoint", "--layer-id", "layer-1", "--checkpoint-id", "cp-1"}, {"fs", "rollback-layer", "--layer-id", "layer-1"}, {"fs", "commit-layer", "--layer-id", "layer-1"}, @@ -300,6 +324,8 @@ func TestLiveFSCommandSurface(t *testing.T) { } testLiveReadOnlyDryRunRejections(t, bin, profileName, [][]string{ {"fs", "check-file-system"}, {"fs", "list-file-systems"}, {"fs", "describe-file-system"}, + {"fs", "describe-file-system-extract-configuration", "--file-system-id", selected.FSTenantID, "--media-type", "image"}, + {"fs", "describe-file-system-embedding-configuration", "--file-system-id", selected.FSTenantID}, {"fs", "read-file"}, {"fs", "list-files"}, {"fs", "describe-file"}, {"fs", "search-file-content"}, {"fs", "find-files"}, {"fs", "list-layers"}, {"fs", "describe-layer"}, {"fs", "diff-layer"}, diff --git a/internal/api/fs/admin_tenant.go b/internal/api/fs/admin_tenant.go index e9fffba..5b0814c 100644 --- a/internal/api/fs/admin_tenant.go +++ b/internal/api/fs/admin_tenant.go @@ -26,13 +26,17 @@ type AdminTenantQuotaConfig struct { MaxStorageSize int64 `json:"max_storage_size"` MaxFileSize int64 `json:"max_file_size"` MaxFileCount int64 `json:"max_file_count"` + MaxMediaLLMFiles int64 `json:"max_media_llm_files"` + MaxVideoLLMFiles int64 `json:"max_video_llm_files"` TiDBCloudSpendingLimit *int64 `json:"tidbcloud_spending_limit"` } type AdminTenantQuotaUsage struct { - StorageBytes int64 `json:"storage_bytes"` - ReservedBytes int64 `json:"reserved_bytes"` - FileCount int64 `json:"file_count"` + StorageBytes int64 `json:"storage_bytes"` + ReservedBytes int64 `json:"reserved_bytes"` + FileCount int64 `json:"file_count"` + MediaFileCount int64 `json:"media_file_count"` + VideoFileCount int64 `json:"video_file_count"` } type AdminTenantQuota struct { diff --git a/internal/api/fs/ai_configuration.go b/internal/api/fs/ai_configuration.go new file mode 100644 index 0000000..e97e259 --- /dev/null +++ b/internal/api/fs/ai_configuration.go @@ -0,0 +1,113 @@ +package fs + +import ( + "context" + "net/http" + "net/url" + "time" +) + +type ExtractMediaType string + +const ( + ExtractMediaTypeImage ExtractMediaType = "image" + ExtractMediaTypeAudio ExtractMediaType = "audio" + ExtractMediaTypeVideo ExtractMediaType = "video" +) + +type AdminTenantExtractConfiguration struct { + Enabled bool `json:"enabled"` + APIBase *string `json:"api_base,omitempty"` + APIKey *string `json:"api_key,omitempty"` + Model *string `json:"model,omitempty"` + Protocol *string `json:"protocol,omitempty"` + Prompt *string `json:"prompt,omitempty"` + Source string `json:"source"` + UpdatedAt *time.Time `json:"updated_at,omitempty"` +} + +type UpdateAdminTenantExtractConfigurationRequest struct { + Enabled *bool `json:"enabled,omitempty"` + APIBase *string `json:"api_base,omitempty"` + APIKey *string `json:"api_key,omitempty"` + Model *string `json:"model,omitempty"` + Protocol *string `json:"protocol,omitempty"` + Prompt *string `json:"prompt,omitempty"` +} + +type AdminTenantEmbeddingConfiguration struct { + Enabled bool `json:"enabled"` + APIBase *string `json:"api_base,omitempty"` + APIKey *string `json:"api_key,omitempty"` + Model *string `json:"model,omitempty"` + Source string `json:"source"` + Generation uint64 `json:"generation,omitempty"` + UpdatedAt *time.Time `json:"updated_at,omitempty"` +} + +type UpdateAdminTenantEmbeddingConfigurationRequest struct { + Enabled bool `json:"enabled"` + APIBase *string `json:"api_base,omitempty"` + APIKey *string `json:"api_key,omitempty"` + Model *string `json:"model,omitempty"` +} + +func (c *Client) GetAdminTenantExtractConfiguration(ctx context.Context, creds TiDBCloudCredentials, fileSystemID string, mediaType ExtractMediaType) (AdminTenantExtractConfiguration, error) { + req, err := c.api.NewRequest(ctx, http.MethodGet, adminTenantExtractConfigurationPath(fileSystemID, mediaType), nil) + if err != nil { + return AdminTenantExtractConfiguration{}, err + } + setTiDBCloudCredentialHeaders(req, creds) + var response AdminTenantExtractConfiguration + if err := c.api.DoJSON(req, &response); err != nil { + return AdminTenantExtractConfiguration{}, err + } + return response, nil +} + +func (c *Client) UpdateAdminTenantExtractConfiguration(ctx context.Context, creds TiDBCloudCredentials, fileSystemID string, mediaType ExtractMediaType, input UpdateAdminTenantExtractConfigurationRequest) (AdminTenantExtractConfiguration, error) { + req, err := c.api.NewRequest(ctx, http.MethodPut, adminTenantExtractConfigurationPath(fileSystemID, mediaType), input) + if err != nil { + return AdminTenantExtractConfiguration{}, err + } + setTiDBCloudCredentialHeaders(req, creds) + var response AdminTenantExtractConfiguration + if err := c.api.DoJSON(req, &response); err != nil { + return AdminTenantExtractConfiguration{}, err + } + return response, nil +} + +func (c *Client) GetAdminTenantEmbeddingConfiguration(ctx context.Context, creds TiDBCloudCredentials, fileSystemID string) (AdminTenantEmbeddingConfiguration, error) { + req, err := c.api.NewRequest(ctx, http.MethodGet, adminTenantEmbeddingConfigurationPath(fileSystemID), nil) + if err != nil { + return AdminTenantEmbeddingConfiguration{}, err + } + setTiDBCloudCredentialHeaders(req, creds) + var response AdminTenantEmbeddingConfiguration + if err := c.api.DoJSON(req, &response); err != nil { + return AdminTenantEmbeddingConfiguration{}, err + } + return response, nil +} + +func (c *Client) UpdateAdminTenantEmbeddingConfiguration(ctx context.Context, creds TiDBCloudCredentials, fileSystemID string, input UpdateAdminTenantEmbeddingConfigurationRequest) (AdminTenantEmbeddingConfiguration, error) { + req, err := c.api.NewRequest(ctx, http.MethodPut, adminTenantEmbeddingConfigurationPath(fileSystemID), input) + if err != nil { + return AdminTenantEmbeddingConfiguration{}, err + } + setTiDBCloudCredentialHeaders(req, creds) + var response AdminTenantEmbeddingConfiguration + if err := c.api.DoJSON(req, &response); err != nil { + return AdminTenantEmbeddingConfiguration{}, err + } + return response, nil +} + +func adminTenantExtractConfigurationPath(fileSystemID string, mediaType ExtractMediaType) string { + return "/v1/admin/tenants/" + url.PathEscape(fileSystemID) + "/extract-config/" + url.PathEscape(string(mediaType)) +} + +func adminTenantEmbeddingConfigurationPath(fileSystemID string) string { + return "/v1/admin/tenants/" + url.PathEscape(fileSystemID) + "/embedding-config" +} diff --git a/internal/api/fs/ai_configuration_test.go b/internal/api/fs/ai_configuration_test.go new file mode 100644 index 0000000..16f6ea3 --- /dev/null +++ b/internal/api/fs/ai_configuration_test.go @@ -0,0 +1,100 @@ +package fs + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" +) + +func TestAdminTenantAIConfigurationClientContracts(t *testing.T) { + const publicKey = "public-secret" + const privateKey = "private-secret" + updatedAt := time.Date(2026, 8, 27, 10, 0, 0, 0, time.UTC) + requests := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requests++ + if r.Header.Get(tidbCloudPublicKeyHeader) != publicKey || r.Header.Get(tidbCloudPrivateKeyHeader) != privateKey { + t.Fatalf("credential headers = %q/%q", r.Header.Get(tidbCloudPublicKeyHeader), r.Header.Get(tidbCloudPrivateKeyHeader)) + } + switch { + case r.Method == http.MethodGet && r.URL.Path == "/v1/admin/tenants/tenant-1/extract-config/audio": + _ = json.NewEncoder(w).Encode(AdminTenantExtractConfiguration{Enabled: false, Source: "none"}) + case r.Method == http.MethodPut && r.URL.Path == "/v1/admin/tenants/tenant-1/extract-config/audio": + var body map[string]any + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Fatal(err) + } + if len(body) != 6 || body["enabled"] != true || body["api_base"] != "https://provider.example/v1" || body["api_key"] != "provider-secret" || body["model"] != "audio-model" || body["protocol"] != "qwen-asr" || body["prompt"] != "transcribe" { + t.Fatalf("extract update body = %#v", body) + } + masked := "pro********" + base, model, protocol, prompt := "https://provider.example/v1", "audio-model", "qwen-asr", "transcribe" + _ = json.NewEncoder(w).Encode(AdminTenantExtractConfiguration{Enabled: true, APIBase: &base, APIKey: &masked, Model: &model, Protocol: &protocol, Prompt: &prompt, Source: "custom", UpdatedAt: &updatedAt}) + case r.Method == http.MethodGet && r.URL.Path == "/v1/admin/tenants/tenant-1/embedding-config": + _ = json.NewEncoder(w).Encode(AdminTenantEmbeddingConfiguration{Enabled: false, Source: "database_auto"}) + case r.Method == http.MethodPut && r.URL.Path == "/v1/admin/tenants/tenant-1/embedding-config": + var body map[string]any + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Fatal(err) + } + if len(body) != 4 || body["enabled"] != true || body["api_base"] != "https://provider.example/v1" || body["api_key"] != "provider-secret" || body["model"] != "embedding-model" { + t.Fatalf("embedding update body = %#v", body) + } + masked := "pro********" + base, model := "https://provider.example/v1", "embedding-model" + _ = json.NewEncoder(w).Encode(AdminTenantEmbeddingConfiguration{Enabled: true, APIBase: &base, APIKey: &masked, Model: &model, Source: "custom", Generation: 3, UpdatedAt: &updatedAt}) + default: + t.Fatalf("unexpected request %s %s", r.Method, r.URL.Path) + } + })) + defer server.Close() + + client := testClient(t, server.URL) + creds := TiDBCloudCredentials{PublicKey: publicKey, PrivateKey: privateKey} + if result, err := client.GetAdminTenantExtractConfiguration(context.Background(), creds, "tenant-1", ExtractMediaTypeAudio); err != nil || result.Source != "none" { + t.Fatalf("GetAdminTenantExtractConfiguration() = %#v, %v", result, err) + } + enabled := true + base, key, audioModel, protocol, prompt := "https://provider.example/v1", "provider-secret", "audio-model", "qwen-asr", "transcribe" + extract, err := client.UpdateAdminTenantExtractConfiguration(context.Background(), creds, "tenant-1", ExtractMediaTypeAudio, UpdateAdminTenantExtractConfigurationRequest{ + Enabled: &enabled, APIBase: &base, APIKey: &key, Model: &audioModel, Protocol: &protocol, Prompt: &prompt, + }) + if err != nil || extract.APIKey == nil || *extract.APIKey != "pro********" { + t.Fatalf("UpdateAdminTenantExtractConfiguration() = %#v, %v", extract, err) + } + if result, err := client.GetAdminTenantEmbeddingConfiguration(context.Background(), creds, "tenant-1"); err != nil || result.Source != "database_auto" { + t.Fatalf("GetAdminTenantEmbeddingConfiguration() = %#v, %v", result, err) + } + embeddingModel := "embedding-model" + embedding, err := client.UpdateAdminTenantEmbeddingConfiguration(context.Background(), creds, "tenant-1", UpdateAdminTenantEmbeddingConfigurationRequest{ + Enabled: true, APIBase: &base, APIKey: &key, Model: &embeddingModel, + }) + if err != nil || embedding.Generation != 3 { + t.Fatalf("UpdateAdminTenantEmbeddingConfiguration() = %#v, %v", embedding, err) + } + if requests != 4 { + t.Fatalf("requests = %d, want 4", requests) + } +} + +func TestAdminTenantAIConfigurationUpdateDoesNotRetry(t *testing.T) { + var requests atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + requests.Add(1) + http.Error(w, `{"error":"provider request failed"}`, http.StatusBadGateway) + })) + defer server.Close() + + enabled := true + _, err := testClient(t, server.URL).UpdateAdminTenantExtractConfiguration(context.Background(), TiDBCloudCredentials{PublicKey: "public", PrivateKey: "private"}, "tenant-1", ExtractMediaTypeImage, UpdateAdminTenantExtractConfigurationRequest{Enabled: &enabled}) + if err == nil { + t.Fatal("update should fail") + } + if got := requests.Load(); got != 1 { + t.Fatalf("requests = %d, want one non-retried PUT", got) + } +} diff --git a/internal/authz/authz.go b/internal/authz/authz.go index a59dd12..aef34c6 100644 --- a/internal/authz/authz.go +++ b/internal/authz/authz.go @@ -10,104 +10,112 @@ import ( type Permission string const ( - DBClusterDiscover Permission = "db.cluster.discover" - StarterClusterRead Permission = "starter.cluster.read" - StarterClusterCreate Permission = "starter.cluster.create" - StarterClusterUpdate Permission = "starter.cluster.update" - StarterClusterDelete Permission = "starter.cluster.delete" - StarterBranchRead Permission = "starter.branch.read" - StarterBranchCreate Permission = "starter.branch.create" - StarterBranchDelete Permission = "starter.branch.delete" - StarterSQLUserRead Permission = "starter.sql_user.read" - StarterSQLUserCreate Permission = "starter.sql_user.create" - StarterSQLUserUpdate Permission = "starter.sql_user.update" - StarterSQLExecute Permission = "starter.sql.execute" - FSVolumeRead Permission = "fs.volume.read" - FSVolumeCreate Permission = "fs.volume.create" - FSVolumeDelete Permission = "fs.volume.delete" - FSTokenList Permission = "fs.token.list" - FSTokenGenerate Permission = "fs.token.generate" - FSTokenIssueScoped Permission = "fs.token.issue_scoped" - FSTokenEnable Permission = "fs.token.enable" - FSTokenDisable Permission = "fs.token.disable" - FSTokenDelete Permission = "fs.token.delete" - FSTokenRefresh Permission = "fs.token.refresh" - FSFileRead Permission = "fs.file.read" - FSFileWrite Permission = "fs.file.write" - FSVaultSecretRead Permission = "fs.vault.secret.read" - FSVaultSecretCreate Permission = "fs.vault.secret.create" - FSVaultSecretUpdate Permission = "fs.vault.secret.update" - FSVaultSecretDelete Permission = "fs.vault.secret.delete" - FSVaultGrantCreate Permission = "fs.vault.grant.create" - FSVaultGrantDelete Permission = "fs.vault.grant.delete" - FSVaultAuditRead Permission = "fs.vault.audit.read" - FSJournalCreate Permission = "fs.journal.create" - FSJournalAppend Permission = "fs.journal.append" - FSJournalRead Permission = "fs.journal.read" - FSJournalSearch Permission = "fs.journal.search" - FSJournalVerify Permission = "fs.journal.verify" - FSGitWorkspaceRead Permission = "fs.git_workspace.read" - FSGitWorkspaceWrite Permission = "fs.git_workspace.write" - FSMount Permission = "fs.mount" + DBClusterDiscover Permission = "db.cluster.discover" + StarterClusterRead Permission = "starter.cluster.read" + StarterClusterCreate Permission = "starter.cluster.create" + StarterClusterUpdate Permission = "starter.cluster.update" + StarterClusterDelete Permission = "starter.cluster.delete" + StarterBranchRead Permission = "starter.branch.read" + StarterBranchCreate Permission = "starter.branch.create" + StarterBranchDelete Permission = "starter.branch.delete" + StarterSQLUserRead Permission = "starter.sql_user.read" + StarterSQLUserCreate Permission = "starter.sql_user.create" + StarterSQLUserUpdate Permission = "starter.sql_user.update" + StarterSQLExecute Permission = "starter.sql.execute" + FSVolumeRead Permission = "fs.volume.read" + FSVolumeCreate Permission = "fs.volume.create" + FSVolumeDelete Permission = "fs.volume.delete" + FSTokenList Permission = "fs.token.list" + FSTokenGenerate Permission = "fs.token.generate" + FSTokenIssueScoped Permission = "fs.token.issue_scoped" + FSTokenEnable Permission = "fs.token.enable" + FSTokenDisable Permission = "fs.token.disable" + FSTokenDelete Permission = "fs.token.delete" + FSTokenRefresh Permission = "fs.token.refresh" + FSExtractConfigRead Permission = "fs.extract_config.read" + FSExtractConfigUpdate Permission = "fs.extract_config.update" + FSEmbeddingConfigRead Permission = "fs.embedding_config.read" + FSEmbeddingConfigUpdate Permission = "fs.embedding_config.update" + FSFileRead Permission = "fs.file.read" + FSFileWrite Permission = "fs.file.write" + FSVaultSecretRead Permission = "fs.vault.secret.read" + FSVaultSecretCreate Permission = "fs.vault.secret.create" + FSVaultSecretUpdate Permission = "fs.vault.secret.update" + FSVaultSecretDelete Permission = "fs.vault.secret.delete" + FSVaultGrantCreate Permission = "fs.vault.grant.create" + FSVaultGrantDelete Permission = "fs.vault.grant.delete" + FSVaultAuditRead Permission = "fs.vault.audit.read" + FSJournalCreate Permission = "fs.journal.create" + FSJournalAppend Permission = "fs.journal.append" + FSJournalRead Permission = "fs.journal.read" + FSJournalSearch Permission = "fs.journal.search" + FSJournalVerify Permission = "fs.journal.verify" + FSGitWorkspaceRead Permission = "fs.git_workspace.read" + FSGitWorkspaceWrite Permission = "fs.git_workspace.write" + FSMount Permission = "fs.mount" ) var commandPermissions = map[string]Permission{ - "ti fs create-file-system": FSVolumeCreate, - "ti fs delete-file-system": FSVolumeDelete, - "ti fs list-file-systems": FSVolumeRead, - "ti fs describe-file-system": FSVolumeRead, - "ti fs check-file-system": FSVolumeRead, - "ti fs generate-file-system-token": FSTokenGenerate, - "ti fs generate-file-system-scoped-token": FSTokenIssueScoped, - "ti fs list-file-system-tokens": FSTokenList, - "ti fs enable-file-system-token": FSTokenEnable, - "ti fs disable-file-system-token": FSTokenDisable, - "ti fs delete-file-system-token": FSTokenDelete, - "ti fs refresh-file-system-token": FSTokenRefresh, - "ti fs copy-file": FSFileWrite, - "ti fs read-file": FSFileRead, - "ti fs list-files": FSFileRead, - "ti fs describe-file": FSFileRead, - "ti fs move-file": FSFileWrite, - "ti fs delete-file": FSFileWrite, - "ti fs create-directory": FSFileWrite, - "ti fs chmod-file": FSFileWrite, - "ti fs create-symlink": FSFileWrite, - "ti fs create-hardlink": FSFileWrite, - "ti fs search-file-content": FSFileRead, - "ti fs find-files": FSFileRead, - "ti fs create-layer": FSFileWrite, - "ti fs list-layers": FSFileRead, - "ti fs describe-layer": FSFileRead, - "ti fs diff-layer": FSFileRead, - "ti fs create-layer-checkpoint": FSFileWrite, - "ti fs rollback-layer": FSFileWrite, - "ti fs commit-layer": FSFileWrite, - "ti fs pack-file-system": FSFileWrite, - "ti fs unpack-file-system": FSFileRead, - "ti fs mount-file-system": FSMount, - "ti fs drain-file-system": FSMount, - "ti fs unmount-file-system": FSMount, - "ti fs-vault create-secret": FSVaultSecretCreate, - "ti fs-vault replace-secret": FSVaultSecretUpdate, - "ti fs-vault read-secret": FSVaultSecretRead, - "ti fs-vault list-secrets": FSVaultSecretRead, - "ti fs-vault delete-secret": FSVaultSecretDelete, - "ti fs-vault create-grant": FSVaultGrantCreate, - "ti fs-vault delete-grant": FSVaultGrantDelete, - "ti fs-vault list-audit-events": FSVaultAuditRead, - "ti fs-vault run-with-secret": FSVaultSecretRead, - "ti fs-vault mount-vault": FSVaultSecretRead, - "ti fs-vault unmount-vault": FSVaultSecretRead, - "ti fs-journal create-journal": FSJournalCreate, - "ti fs-journal append-journal-entries": FSJournalAppend, - "ti fs-journal read-journal-entries": FSJournalRead, - "ti fs-journal search-journal-entries": FSJournalSearch, - "ti fs-journal verify-journal": FSJournalVerify, - "ti fs-git clone-git-workspace": FSGitWorkspaceWrite, - "ti fs-git hydrate-git-workspace": FSGitWorkspaceRead, - "ti fs-git add-git-worktree": FSGitWorkspaceWrite, - "ti fs-git remove-git-worktree": FSGitWorkspaceWrite, + "ti fs create-file-system": FSVolumeCreate, + "ti fs delete-file-system": FSVolumeDelete, + "ti fs list-file-systems": FSVolumeRead, + "ti fs describe-file-system": FSVolumeRead, + "ti fs check-file-system": FSVolumeRead, + "ti fs generate-file-system-token": FSTokenGenerate, + "ti fs generate-file-system-scoped-token": FSTokenIssueScoped, + "ti fs list-file-system-tokens": FSTokenList, + "ti fs enable-file-system-token": FSTokenEnable, + "ti fs disable-file-system-token": FSTokenDisable, + "ti fs delete-file-system-token": FSTokenDelete, + "ti fs refresh-file-system-token": FSTokenRefresh, + "ti fs describe-file-system-extract-configuration": FSExtractConfigRead, + "ti fs update-file-system-extract-configuration": FSExtractConfigUpdate, + "ti fs describe-file-system-embedding-configuration": FSEmbeddingConfigRead, + "ti fs update-file-system-embedding-configuration": FSEmbeddingConfigUpdate, + "ti fs copy-file": FSFileWrite, + "ti fs read-file": FSFileRead, + "ti fs list-files": FSFileRead, + "ti fs describe-file": FSFileRead, + "ti fs move-file": FSFileWrite, + "ti fs delete-file": FSFileWrite, + "ti fs create-directory": FSFileWrite, + "ti fs chmod-file": FSFileWrite, + "ti fs create-symlink": FSFileWrite, + "ti fs create-hardlink": FSFileWrite, + "ti fs search-file-content": FSFileRead, + "ti fs find-files": FSFileRead, + "ti fs create-layer": FSFileWrite, + "ti fs list-layers": FSFileRead, + "ti fs describe-layer": FSFileRead, + "ti fs diff-layer": FSFileRead, + "ti fs create-layer-checkpoint": FSFileWrite, + "ti fs rollback-layer": FSFileWrite, + "ti fs commit-layer": FSFileWrite, + "ti fs pack-file-system": FSFileWrite, + "ti fs unpack-file-system": FSFileRead, + "ti fs mount-file-system": FSMount, + "ti fs drain-file-system": FSMount, + "ti fs unmount-file-system": FSMount, + "ti fs-vault create-secret": FSVaultSecretCreate, + "ti fs-vault replace-secret": FSVaultSecretUpdate, + "ti fs-vault read-secret": FSVaultSecretRead, + "ti fs-vault list-secrets": FSVaultSecretRead, + "ti fs-vault delete-secret": FSVaultSecretDelete, + "ti fs-vault create-grant": FSVaultGrantCreate, + "ti fs-vault delete-grant": FSVaultGrantDelete, + "ti fs-vault list-audit-events": FSVaultAuditRead, + "ti fs-vault run-with-secret": FSVaultSecretRead, + "ti fs-vault mount-vault": FSVaultSecretRead, + "ti fs-vault unmount-vault": FSVaultSecretRead, + "ti fs-journal create-journal": FSJournalCreate, + "ti fs-journal append-journal-entries": FSJournalAppend, + "ti fs-journal read-journal-entries": FSJournalRead, + "ti fs-journal search-journal-entries": FSJournalSearch, + "ti fs-journal verify-journal": FSJournalVerify, + "ti fs-git clone-git-workspace": FSGitWorkspaceWrite, + "ti fs-git hydrate-git-workspace": FSGitWorkspaceRead, + "ti fs-git add-git-worktree": FSGitWorkspaceWrite, + "ti fs-git remove-git-worktree": FSGitWorkspaceWrite, } func ForCommand(commandPath string) (Permission, error) { diff --git a/internal/cli/commands.go b/internal/cli/commands.go index 79049be..3e1d376 100644 --- a/internal/cli/commands.go +++ b/internal/cli/commands.go @@ -20,6 +20,7 @@ import ( dbstarter "github.com/tidbcloud/ti-cli/internal/db/product/starter" "github.com/tidbcloud/ti-cli/internal/dryrun" tifs "github.com/tidbcloud/ti-cli/internal/fs" + "github.com/tidbcloud/ti-cli/internal/fs/aiconfig" "github.com/tidbcloud/ti-cli/internal/fs/fscred" "github.com/tidbcloud/ti-cli/internal/fs/tokenmgmt" outputpkg "github.com/tidbcloud/ti-cli/internal/output" @@ -859,6 +860,10 @@ func newFSCommand(info version.Info) *cobra.Command { newFSDeleteFileSystemCommand(info), newFSListFileSystemsCommand(info), newFSDescribeFileSystemCommand(info), + newFSDescribeFileSystemExtractConfigurationCommand(info), + newFSUpdateFileSystemExtractConfigurationCommand(info), + newFSDescribeFileSystemEmbeddingConfigurationCommand(info), + newFSUpdateFileSystemEmbeddingConfigurationCommand(info), newFSImportFileSystemTokenCommand(info), newFSGenerateFileSystemTokenCommand(info), newFSGenerateFileSystemScopedTokenCommand(info), @@ -894,7 +899,11 @@ func newFSCommand(info version.Info) *cobra.Command { newFSUnmountFileSystemCommand(info), } tokenCommands := []string{"generate-file-system-token", "generate-file-system-scoped-token", "list-file-system-tokens", "enable-file-system-token", "disable-file-system-token", "delete-file-system-token", "refresh-file-system-token"} - selectorExclusions := append([]string{"create-file-system", "list-file-systems", "describe-file-system", "delete-file-system", "import-file-system-token", "drain-file-system", "unmount-file-system"}, tokenCommands...) + selectorExclusions := append([]string{ + "create-file-system", "list-file-systems", "describe-file-system", "delete-file-system", "import-file-system-token", "drain-file-system", "unmount-file-system", + "describe-file-system-extract-configuration", "update-file-system-extract-configuration", + "describe-file-system-embedding-configuration", "update-file-system-embedding-configuration", + }, tokenCommands...) addFSSelectorFlags(commands, selectorExclusions...) addFSAuthFlags(commands, "create-file-system", @@ -911,6 +920,10 @@ func newFSCommand(info version.Info) *cobra.Command { "disable-file-system-token", "delete-file-system-token", "refresh-file-system-token", + "describe-file-system-extract-configuration", + "update-file-system-extract-configuration", + "describe-file-system-embedding-configuration", + "update-file-system-embedding-configuration", ) cmd.AddCommand(commands...) return cmd @@ -1369,6 +1382,134 @@ func newFSDescribeFileSystemCommand(info version.Info) *cobra.Command { return cmd } +func newFSDescribeFileSystemExtractConfigurationCommand(info version.Info) *cobra.Command { + cmd := newControlPlaneCommand(controlPlaneCommandSpec{ + Use: "describe-file-system-extract-configuration", + Short: "Describe media extraction configuration for a file system.", + Mutation: readOnlyCommand, + Permission: authz.FSExtractConfigRead, + Run: func(ctx commandContext) (any, error) { + service, profile, err := fsAIConfigurationServiceAndProfile(ctx) + if err != nil { + return nil, err + } + opts, err := fsDescribeExtractConfigurationOptions(ctx, profile) + if err != nil { + return nil, err + } + return service.DescribeExtract(ctx.cmd.Context(), opts) + }, + }, info) + cmd.Flags().String("file-system-id", "", "The file system ID. This command requires TiDB Cloud API credentials.") + cmd.Flags().String("media-type", "", "Media type: image, audio, or video.") + markUsageRequired(cmd, "file-system-id", "media-type") + return cmd +} + +func newFSUpdateFileSystemExtractConfigurationCommand(info version.Info) *cobra.Command { + cmd := newControlPlaneCommand(controlPlaneCommandSpec{ + Use: "update-file-system-extract-configuration", + Short: "Update media extraction configuration for a file system.", + Long: "Update media extraction configuration for a file system. Enabling or replacing a provider performs a real provider validation request that can incur a charge. If an update outcome is ambiguous, describe the configuration before retrying.", + Mutation: mutatingCommand, + Permission: authz.FSExtractConfigUpdate, + Run: func(ctx commandContext) (any, error) { + service, profile, err := fsAIConfigurationServiceAndProfile(ctx) + if err != nil { + return nil, err + } + opts, err := fsUpdateExtractConfigurationOptions(ctx, profile) + if err != nil { + return nil, err + } + return service.UpdateExtract(ctx.cmd.Context(), opts) + }, + DryRun: func(ctx commandContext) (dryrun.Result, error) { + service, profile, err := fsAIConfigurationServiceAndProfile(ctx) + if err != nil { + return dryrun.Result{}, err + } + opts, err := fsUpdateExtractConfigurationOptions(ctx, profile) + if err != nil { + return dryrun.Result{}, err + } + return service.DryRunUpdateExtract(ctx.CommandPath(), opts) + }, + }, info) + cmd.Flags().String("file-system-id", "", "The file system ID. This command requires TiDB Cloud API credentials.") + cmd.Flags().String("media-type", "", "Media type: image, audio, or video.") + cmd.Flags().Bool("enabled", false, "Enable or disable extraction. An explicit true or false value is required when this option is used.") + requireExplicitBoolValue(cmd, "enabled") + cmd.Flags().String("provider-api-base", "", "HTTPS base URL for the extraction provider API.") + cmd.Flags().String("provider-model", "", "Model name accepted by the extraction provider.") + cmd.Flags().String("provider-protocol", "", "Provider protocol: openai, or qwen-asr for audio.") + cmd.Flags().String("prompt", "", "Extraction prompt. Pass an empty value to restore the backend default prompt.") + markUsageRequired(cmd, "file-system-id", "media-type") + return cmd +} + +func newFSDescribeFileSystemEmbeddingConfigurationCommand(info version.Info) *cobra.Command { + cmd := newControlPlaneCommand(controlPlaneCommandSpec{ + Use: "describe-file-system-embedding-configuration", + Short: "Describe embedding configuration for a file system.", + Mutation: readOnlyCommand, + Permission: authz.FSEmbeddingConfigRead, + Run: func(ctx commandContext) (any, error) { + service, profile, err := fsAIConfigurationServiceAndProfile(ctx) + if err != nil { + return nil, err + } + opts, err := fsDescribeEmbeddingConfigurationOptions(ctx, profile) + if err != nil { + return nil, err + } + return service.DescribeEmbedding(ctx.cmd.Context(), opts) + }, + }, info) + cmd.Flags().String("file-system-id", "", "The file system ID. This command requires TiDB Cloud API credentials.") + markUsageRequired(cmd, "file-system-id") + return cmd +} + +func newFSUpdateFileSystemEmbeddingConfigurationCommand(info version.Info) *cobra.Command { + cmd := newControlPlaneCommand(controlPlaneCommandSpec{ + Use: "update-file-system-embedding-configuration", + Short: "Replace embedding configuration for a file system.", + Long: "Replace embedding configuration for a file system. Enabling a provider performs a real provider validation request that can incur a charge. If an update outcome is ambiguous, describe the configuration before retrying.", + Mutation: mutatingCommand, + Permission: authz.FSEmbeddingConfigUpdate, + Run: func(ctx commandContext) (any, error) { + service, profile, err := fsAIConfigurationServiceAndProfile(ctx) + if err != nil { + return nil, err + } + opts, err := fsUpdateEmbeddingConfigurationOptions(ctx, profile) + if err != nil { + return nil, err + } + return service.UpdateEmbedding(ctx.cmd.Context(), opts) + }, + DryRun: func(ctx commandContext) (dryrun.Result, error) { + service, profile, err := fsAIConfigurationServiceAndProfile(ctx) + if err != nil { + return dryrun.Result{}, err + } + opts, err := fsUpdateEmbeddingConfigurationOptions(ctx, profile) + if err != nil { + return dryrun.Result{}, err + } + return service.DryRunUpdateEmbedding(ctx.CommandPath(), opts) + }, + }, info) + cmd.Flags().String("file-system-id", "", "The file system ID. This command requires TiDB Cloud API credentials.") + cmd.Flags().Bool("enabled", false, "Enable or disable app-managed embedding. An explicit true or false value is required.") + requireExplicitBoolValue(cmd, "enabled") + cmd.Flags().String("provider-api-base", "", "HTTPS base URL for an OpenAI-compatible embedding API.") + cmd.Flags().String("provider-model", "", "Embedding model name. The provider must return 1024 dimensions.") + markUsageRequired(cmd, "file-system-id", "enabled") + return cmd +} + func newFSDeleteFileSystemCommand(info version.Info) *cobra.Command { cmd := newControlPlaneCommand(controlPlaneCommandSpec{ Use: "delete-file-system", @@ -2795,6 +2936,130 @@ func fsService(ctx commandContext, profile *config.Profile) (tifs.Service, *conf return service, profile, nil } +func fsAIConfigurationServiceAndProfile(ctx commandContext) (aiconfig.Service, *config.Profile, error) { + profile, err := ctx.LoadProfile() + if err != nil { + return aiconfig.Service{}, nil, err + } + debug, err := ctx.BoolFlag("debug") + if err != nil { + return aiconfig.Service{}, nil, err + } + return aiconfig.Service{Timeout: 30 * time.Second, Debug: debug, DebugWriter: ctx.cmd.ErrOrStderr()}, profile, nil +} + +func fsDescribeExtractConfigurationOptions(ctx commandContext, profile *config.Profile) (aiconfig.DescribeExtractOptions, error) { + fileSystemID, err := ctx.StringFlag("file-system-id") + if err != nil { + return aiconfig.DescribeExtractOptions{}, err + } + mediaType, err := ctx.StringFlag("media-type") + if err != nil { + return aiconfig.DescribeExtractOptions{}, err + } + return aiconfig.DescribeExtractOptions{Profile: profile, FileSystemID: fileSystemID, MediaType: mediaType}, nil +} + +func fsUpdateExtractConfigurationOptions(ctx commandContext, profile *config.Profile) (aiconfig.UpdateExtractOptions, error) { + describe, err := fsDescribeExtractConfigurationOptions(ctx, profile) + if err != nil { + return aiconfig.UpdateExtractOptions{}, err + } + enabled, err := optionalBoolFlag(ctx, "enabled") + if err != nil { + return aiconfig.UpdateExtractOptions{}, err + } + apiBase, err := optionalStringFlag(ctx, "provider-api-base") + if err != nil { + return aiconfig.UpdateExtractOptions{}, err + } + model, err := optionalStringFlag(ctx, "provider-model") + if err != nil { + return aiconfig.UpdateExtractOptions{}, err + } + protocol, err := optionalStringFlag(ctx, "provider-protocol") + if err != nil { + return aiconfig.UpdateExtractOptions{}, err + } + prompt, err := optionalStringFlag(ctx, "prompt") + if err != nil { + return aiconfig.UpdateExtractOptions{}, err + } + providerSet := apiBase != nil || model != nil || (enabled != nil && *enabled) || (protocol != nil && strings.EqualFold(strings.TrimSpace(describe.MediaType), aiconfig.MediaTypeAudio)) + providerAPIKey, providerAPIKeySupplied := "", false + if providerSet { + providerAPIKey, providerAPIKeySupplied = os.LookupEnv("TI_FS_AI_PROVIDER_API_KEY") + } + return aiconfig.UpdateExtractOptions{ + Profile: profile, FileSystemID: describe.FileSystemID, MediaType: describe.MediaType, Enabled: enabled, + ProviderAPIBase: apiBase, ProviderModel: model, ProviderProtocol: protocol, Prompt: prompt, + ProviderAPIKey: providerAPIKey, ProviderAPIKeySupplied: providerAPIKeySupplied, + }, nil +} + +func fsDescribeEmbeddingConfigurationOptions(ctx commandContext, profile *config.Profile) (aiconfig.DescribeEmbeddingOptions, error) { + fileSystemID, err := ctx.StringFlag("file-system-id") + if err != nil { + return aiconfig.DescribeEmbeddingOptions{}, err + } + return aiconfig.DescribeEmbeddingOptions{Profile: profile, FileSystemID: fileSystemID}, nil +} + +func fsUpdateEmbeddingConfigurationOptions(ctx commandContext, profile *config.Profile) (aiconfig.UpdateEmbeddingOptions, error) { + describe, err := fsDescribeEmbeddingConfigurationOptions(ctx, profile) + if err != nil { + return aiconfig.UpdateEmbeddingOptions{}, err + } + enabled, err := ctx.BoolFlag("enabled") + if err != nil { + return aiconfig.UpdateEmbeddingOptions{}, err + } + apiBase, err := optionalStringFlag(ctx, "provider-api-base") + if err != nil { + return aiconfig.UpdateEmbeddingOptions{}, err + } + model, err := optionalStringFlag(ctx, "provider-model") + if err != nil { + return aiconfig.UpdateEmbeddingOptions{}, err + } + providerAPIKey, providerAPIKeySupplied := "", false + if enabled { + providerAPIKey, providerAPIKeySupplied = os.LookupEnv("TI_FS_AI_PROVIDER_API_KEY") + } + return aiconfig.UpdateEmbeddingOptions{ + Profile: profile, FileSystemID: describe.FileSystemID, Enabled: enabled, + ProviderAPIBase: apiBase, ProviderModel: model, ProviderAPIKey: providerAPIKey, ProviderAPIKeySupplied: providerAPIKeySupplied, + }, nil +} + +func optionalStringFlag(ctx commandContext, name string) (*string, error) { + if !ctx.FlagChanged(name) { + return nil, nil + } + value, err := ctx.StringFlag(name) + if err != nil { + return nil, err + } + return &value, nil +} + +func optionalBoolFlag(ctx commandContext, name string) (*bool, error) { + if !ctx.FlagChanged(name) { + return nil, nil + } + value, err := ctx.BoolFlag(name) + if err != nil { + return nil, err + } + return &value, nil +} + +func requireExplicitBoolValue(cmd *cobra.Command, name string) { + if flag := cmd.Flags().Lookup(name); flag != nil { + flag.NoOptDefVal = "" + } +} + func fsServiceAndProfile(ctx commandContext) (tifs.Service, *config.Profile, error) { return fsAuthenticatedServiceAndProfile(ctx, true) } diff --git a/internal/cli/root.go b/internal/cli/root.go index 6f5ed5d..ed788b1 100644 --- a/internal/cli/root.go +++ b/internal/cli/root.go @@ -949,6 +949,9 @@ func usageFlagValueType(flag *pflag.Flag) string { } switch flag.Value.Type() { case "bool": + if flag.NoOptDefVal == "" { + return "boolean" + } return "" case "stringArray": return "string" diff --git a/internal/cli/root_test.go b/internal/cli/root_test.go index fd37bcd..1ef5034 100644 --- a/internal/cli/root_test.go +++ b/internal/cli/root_test.go @@ -883,13 +883,17 @@ func TestFSOperationalCommandsExposeResourceSelector(t *testing.T) { func TestFSRemoteCommandsExposeTokenFlag(t *testing.T) { root := NewRootCommand(testVersion()) excluded := map[string]bool{ - "ti fs create-file-system": true, - "ti fs list-file-systems": true, - "ti fs describe-file-system": true, - "ti fs delete-file-system": true, - "ti fs drain-file-system": true, - "ti fs unmount-file-system": true, - "ti fs-vault unmount-vault": true, + "ti fs create-file-system": true, + "ti fs list-file-systems": true, + "ti fs describe-file-system": true, + "ti fs delete-file-system": true, + "ti fs describe-file-system-extract-configuration": true, + "ti fs update-file-system-extract-configuration": true, + "ti fs describe-file-system-embedding-configuration": true, + "ti fs update-file-system-embedding-configuration": true, + "ti fs drain-file-system": true, + "ti fs unmount-file-system": true, + "ti fs-vault unmount-vault": true, } visitCommands(root, func(cmd *cobra.Command) { if cmd.Name() == "help" || cmd.HasSubCommands() || excluded[cmd.CommandPath()] { diff --git a/internal/fs/aiconfig/output.go b/internal/fs/aiconfig/output.go new file mode 100644 index 0000000..bfe2b71 --- /dev/null +++ b/internal/fs/aiconfig/output.go @@ -0,0 +1,56 @@ +package aiconfig + +import ( + "fmt" + "strings" + "time" +) + +func (r ExtractResult) Human() string { + return strings.Join([]string{ + "File system ID: " + r.FileSystemID, + "Media type: " + r.MediaType, + fmt.Sprintf("Enabled: %t", r.Enabled), + "Source: " + valueOrNone(r.Source), + "Provider API base: " + pointerOrNone(r.APIBase), + "Provider API key: " + pointerOrNone(r.APIKey), + "Provider model: " + pointerOrNone(r.Model), + "Provider protocol: " + pointerOrNone(r.Protocol), + "Prompt: " + pointerOrNone(r.Prompt), + "Updated at: " + timeOrNone(r.UpdatedAt), + }, "\n") +} + +func (r EmbeddingResult) Human() string { + return strings.Join([]string{ + "File system ID: " + r.FileSystemID, + fmt.Sprintf("Enabled: %t", r.Enabled), + "Source: " + valueOrNone(r.Source), + "Provider API base: " + pointerOrNone(r.APIBase), + "Provider API key: " + pointerOrNone(r.APIKey), + "Provider model: " + pointerOrNone(r.Model), + fmt.Sprintf("Generation: %d", r.Generation), + "Updated at: " + timeOrNone(r.UpdatedAt), + }, "\n") +} + +func pointerOrNone(value *string) string { + if value == nil || *value == "" { + return "none" + } + return *value +} + +func valueOrNone(value string) string { + if value == "" { + return "none" + } + return value +} + +func timeOrNone(value *time.Time) string { + if value == nil { + return "none" + } + return value.Format(time.RFC3339) +} diff --git a/internal/fs/aiconfig/service.go b/internal/fs/aiconfig/service.go new file mode 100644 index 0000000..8e9e08b --- /dev/null +++ b/internal/fs/aiconfig/service.go @@ -0,0 +1,479 @@ +package aiconfig + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "strings" + "time" + "unicode/utf8" + + "github.com/tidbcloud/ti-cli/internal/api" + "github.com/tidbcloud/ti-cli/internal/api/endpoints" + apifs "github.com/tidbcloud/ti-cli/internal/api/fs" + apitransport "github.com/tidbcloud/ti-cli/internal/api/transport" + "github.com/tidbcloud/ti-cli/internal/apperr" + "github.com/tidbcloud/ti-cli/internal/auth" + "github.com/tidbcloud/ti-cli/internal/authz" + "github.com/tidbcloud/ti-cli/internal/config" + "github.com/tidbcloud/ti-cli/internal/dryrun" + "github.com/tidbcloud/ti-cli/internal/fs/fscred" +) + +const ( + MediaTypeImage = "image" + MediaTypeAudio = "audio" + MediaTypeVideo = "video" + + ProtocolOpenAI = "openai" + ProtocolQwenASR = "qwen-asr" + + maxPromptBytes = 8 * 1024 +) + +type Service struct { + Resolver endpoints.Resolver + HTTPClient *http.Client + Transport http.RoundTripper + Timeout time.Duration + Debug bool + DebugWriter io.Writer +} + +type DescribeExtractOptions struct { + Profile *config.Profile + FileSystemID string + MediaType string +} + +type UpdateExtractOptions struct { + Profile *config.Profile + FileSystemID string + MediaType string + Enabled *bool + ProviderAPIBase *string + ProviderModel *string + ProviderProtocol *string + Prompt *string + ProviderAPIKey string + ProviderAPIKeySupplied bool +} + +type DescribeEmbeddingOptions struct { + Profile *config.Profile + FileSystemID string +} + +type UpdateEmbeddingOptions struct { + Profile *config.Profile + FileSystemID string + Enabled bool + ProviderAPIBase *string + ProviderModel *string + ProviderAPIKey string + ProviderAPIKeySupplied bool +} + +type ExtractResult struct { + FileSystemID string `json:"file_system_id"` + MediaType string `json:"media_type"` + Enabled bool `json:"enabled"` + APIBase *string `json:"api_base,omitempty"` + APIKey *string `json:"api_key,omitempty"` + Model *string `json:"model,omitempty"` + Protocol *string `json:"protocol,omitempty"` + Prompt *string `json:"prompt,omitempty"` + Source string `json:"source"` + UpdatedAt *time.Time `json:"updated_at,omitempty"` +} + +type EmbeddingResult struct { + FileSystemID string `json:"file_system_id"` + Enabled bool `json:"enabled"` + APIBase *string `json:"api_base,omitempty"` + APIKey *string `json:"api_key,omitempty"` + Model *string `json:"model,omitempty"` + Source string `json:"source"` + Generation uint64 `json:"generation,omitempty"` + UpdatedAt *time.Time `json:"updated_at,omitempty"` +} + +func (s Service) DescribeExtract(ctx context.Context, opts DescribeExtractOptions) (ExtractResult, error) { + id, mediaType, client, creds, err := s.extractInputs(opts.Profile, opts.FileSystemID, opts.MediaType, authz.FSExtractConfigRead, "describe file system extract configuration", "") + if err != nil { + return ExtractResult{}, err + } + response, err := client.GetAdminTenantExtractConfiguration(ctx, creds, id, apifs.ExtractMediaType(mediaType)) + if err != nil { + return ExtractResult{}, mapError(err, "describe", "extract", id, mediaType) + } + if err := validateResponse(response.Source, response.APIKey); err != nil { + return ExtractResult{}, err + } + return extractResult(id, mediaType, response), nil +} + +func (s Service) UpdateExtract(ctx context.Context, opts UpdateExtractOptions) (ExtractResult, error) { + request, err := validateExtractUpdate(opts) + if err != nil { + return ExtractResult{}, err + } + id, mediaType, client, creds, err := s.extractInputs(opts.Profile, opts.FileSystemID, opts.MediaType, authz.FSExtractConfigUpdate, "update file system extract configuration", opts.ProviderAPIKey) + if err != nil { + return ExtractResult{}, err + } + response, err := client.UpdateAdminTenantExtractConfiguration(ctx, creds, id, apifs.ExtractMediaType(mediaType), request) + if err != nil { + return ExtractResult{}, mapError(err, "update", "extract", id, mediaType) + } + if err := validateResponse(response.Source, response.APIKey); err != nil { + return ExtractResult{}, err + } + return extractResult(id, mediaType, response), nil +} + +func (s Service) DescribeEmbedding(ctx context.Context, opts DescribeEmbeddingOptions) (EmbeddingResult, error) { + id, client, creds, err := s.embeddingInputs(opts.Profile, opts.FileSystemID, authz.FSEmbeddingConfigRead, "describe file system embedding configuration", "") + if err != nil { + return EmbeddingResult{}, err + } + response, err := client.GetAdminTenantEmbeddingConfiguration(ctx, creds, id) + if err != nil { + return EmbeddingResult{}, mapError(err, "describe", "embedding", id, "") + } + if err := validateResponse(response.Source, response.APIKey); err != nil { + return EmbeddingResult{}, err + } + return embeddingResult(id, response), nil +} + +func (s Service) UpdateEmbedding(ctx context.Context, opts UpdateEmbeddingOptions) (EmbeddingResult, error) { + request, err := validateEmbeddingUpdate(opts) + if err != nil { + return EmbeddingResult{}, err + } + id, client, creds, err := s.embeddingInputs(opts.Profile, opts.FileSystemID, authz.FSEmbeddingConfigUpdate, "update file system embedding configuration", opts.ProviderAPIKey) + if err != nil { + return EmbeddingResult{}, err + } + response, err := client.UpdateAdminTenantEmbeddingConfiguration(ctx, creds, id, request) + if err != nil { + return EmbeddingResult{}, mapError(err, "update", "embedding", id, "") + } + if err := validateResponse(response.Source, response.APIKey); err != nil { + return EmbeddingResult{}, err + } + return embeddingResult(id, response), nil +} + +func (s Service) DryRunUpdateExtract(command string, opts UpdateExtractOptions) (dryrun.Result, error) { + request, err := validateExtractUpdate(opts) + if err != nil { + return dryrun.Result{}, err + } + id, mediaType, _, _, err := s.extractInputs(opts.Profile, opts.FileSystemID, opts.MediaType, authz.FSExtractConfigUpdate, "update file system extract configuration", opts.ProviderAPIKey) + if err != nil { + return dryrun.Result{}, err + } + body := extractDryRunBody(request, opts.ProviderAPIKeySupplied) + checks := configurationDryRunChecks(opts.Profile, authz.FSExtractConfigUpdate) + checks = append(checks, dryrun.Check{Name: "provider_validation", Status: "skipped", Message: "dry-run does not contact the configured provider"}) + return dryrun.New(command, "update_file_system_extract_configuration", dryrun.RequestSummary{ + Method: http.MethodPut, + Path: "/v1/admin/tenants/" + url.PathEscape(id) + "/extract-config/" + url.PathEscape(mediaType), + Body: body, + }, checks...), nil +} + +func (s Service) DryRunUpdateEmbedding(command string, opts UpdateEmbeddingOptions) (dryrun.Result, error) { + request, err := validateEmbeddingUpdate(opts) + if err != nil { + return dryrun.Result{}, err + } + id, _, _, err := s.embeddingInputs(opts.Profile, opts.FileSystemID, authz.FSEmbeddingConfigUpdate, "update file system embedding configuration", opts.ProviderAPIKey) + if err != nil { + return dryrun.Result{}, err + } + body := embeddingDryRunBody(request, opts.ProviderAPIKeySupplied) + checks := configurationDryRunChecks(opts.Profile, authz.FSEmbeddingConfigUpdate) + checks = append(checks, dryrun.Check{Name: "provider_validation", Status: "skipped", Message: "dry-run does not contact the configured provider"}) + return dryrun.New(command, "update_file_system_embedding_configuration", dryrun.RequestSummary{ + Method: http.MethodPut, + Path: "/v1/admin/tenants/" + url.PathEscape(id) + "/embedding-config", + Body: body, + }, checks...), nil +} + +func configurationDryRunChecks(profile *config.Profile, permission authz.Permission) []dryrun.Check { + return []dryrun.Check{ + {Name: "config_and_credentials", Status: "passed", Message: fmt.Sprintf("profile %q loaded", profile.Name)}, + {Name: "endpoint_selection", Status: "passed", Message: profile.PlacementRegionCode}, + {Name: "permission_requirement", Status: "passed", Message: string(permission)}, + } +} + +func (s Service) extractInputs(profile *config.Profile, fileSystemID, mediaType string, permission authz.Permission, action, providerAPIKey string) (string, string, *apifs.Client, apifs.TiDBCloudCredentials, error) { + id, err := validateFileSystemID(fileSystemID) + if err != nil { + return "", "", nil, apifs.TiDBCloudCredentials{}, err + } + mediaType, err = validateMediaType(mediaType) + if err != nil { + return "", "", nil, apifs.TiDBCloudCredentials{}, err + } + client, creds, err := s.adminClient(profile, permission, action, providerAPIKey) + return id, mediaType, client, creds, err +} + +func (s Service) embeddingInputs(profile *config.Profile, fileSystemID string, permission authz.Permission, action, providerAPIKey string) (string, *apifs.Client, apifs.TiDBCloudCredentials, error) { + id, err := validateFileSystemID(fileSystemID) + if err != nil { + return "", nil, apifs.TiDBCloudCredentials{}, err + } + client, creds, err := s.adminClient(profile, permission, action, providerAPIKey) + return id, client, creds, err +} + +func (s Service) adminClient(profile *config.Profile, permission authz.Permission, action, providerAPIKey string) (*apifs.Client, apifs.TiDBCloudCredentials, error) { + creds, err := auth.ValidateProfile(profile) + if err != nil { + return nil, apifs.TiDBCloudCredentials{}, err + } + provider := profile.FSCloudProvider + regionCode := profile.FSRegionCode + if provider == "" { + provider = profile.CloudProvider + } + if regionCode == "" { + regionCode = profile.RegionCode + } + resolver := s.Resolver + if resolver.IsZero() { + resolver = endpoints.NewResolver() + } + endpoint, err := resolver.ResolveFS(provider, regionCode) + if err != nil { + return nil, apifs.TiDBCloudCredentials{}, err + } + raw, err := api.New(api.Options{ + Endpoint: endpoint, ProfileName: creds.ProfileName, Permission: permission, Action: action, + HTTPClient: s.HTTPClient, Transport: s.Transport, Timeout: s.Timeout, Debug: s.Debug, DebugWriter: s.DebugWriter, + Redactor: apitransport.Redactor{Secrets: []string{creds.PublicKey, creds.PrivateKey, providerAPIKey}}, UserAgent: "ti fs AI configuration", + }) + if err != nil { + return nil, apifs.TiDBCloudCredentials{}, err + } + return apifs.New(raw), apifs.TiDBCloudCredentials{PublicKey: creds.PublicKey, PrivateKey: creds.PrivateKey}, nil +} + +func validateExtractUpdate(opts UpdateExtractOptions) (apifs.UpdateAdminTenantExtractConfigurationRequest, error) { + if opts.Enabled == nil && opts.ProviderAPIBase == nil && opts.ProviderModel == nil && opts.ProviderProtocol == nil && opts.Prompt == nil { + return apifs.UpdateAdminTenantExtractConfigurationRequest{}, usageError("fs.extract_configuration_no_changes", "provide at least one extract configuration option") + } + mediaType, err := validateMediaType(opts.MediaType) + if err != nil { + return apifs.UpdateAdminTenantExtractConfigurationRequest{}, err + } + if opts.Prompt != nil { + if !utf8.ValidString(*opts.Prompt) || len(*opts.Prompt) > maxPromptBytes { + return apifs.UpdateAdminTenantExtractConfigurationRequest{}, usageError("fs.invalid_extract_prompt", fmt.Sprintf("--prompt must be valid UTF-8 and at most %d bytes", maxPromptBytes)) + } + } + if opts.Enabled != nil && !*opts.Enabled { + if opts.ProviderAPIBase != nil || opts.ProviderModel != nil || opts.ProviderProtocol != nil || opts.Prompt != nil { + return apifs.UpdateAdminTenantExtractConfigurationRequest{}, usageError("fs.invalid_extract_configuration", "--enabled false cannot be combined with provider or prompt options") + } + return apifs.UpdateAdminTenantExtractConfigurationRequest{Enabled: opts.Enabled}, nil + } + + providerChange := opts.ProviderAPIBase != nil || opts.ProviderModel != nil || (opts.Enabled != nil && *opts.Enabled) + protocol := opts.ProviderProtocol + if providerChange && protocol == nil { + value := ProtocolOpenAI + protocol = &value + } + if protocol != nil { + normalized := strings.ToLower(strings.TrimSpace(*protocol)) + if normalized == "" { + return apifs.UpdateAdminTenantExtractConfigurationRequest{}, usageError("fs.invalid_provider_protocol", "--provider-protocol must be openai, or qwen-asr for audio") + } + if normalized != ProtocolOpenAI && !(mediaType == MediaTypeAudio && normalized == ProtocolQwenASR) { + return apifs.UpdateAdminTenantExtractConfigurationRequest{}, usageError("fs.invalid_provider_protocol", "--provider-protocol must be openai, or qwen-asr for audio") + } + protocol = &normalized + if mediaType == MediaTypeAudio && opts.ProviderProtocol != nil { + providerChange = true + } + } + if providerChange { + if opts.ProviderAPIBase == nil || opts.ProviderModel == nil || !opts.ProviderAPIKeySupplied { + return apifs.UpdateAdminTenantExtractConfigurationRequest{}, providerTrioError() + } + apiBase := strings.TrimSpace(*opts.ProviderAPIBase) + model := strings.TrimSpace(*opts.ProviderModel) + apiKey := strings.TrimSpace(opts.ProviderAPIKey) + if err := validateProviderFields(apiBase, model, apiKey); err != nil { + return apifs.UpdateAdminTenantExtractConfigurationRequest{}, err + } + opts.ProviderAPIBase = &apiBase + opts.ProviderModel = &model + opts.ProviderAPIKey = apiKey + } + request := apifs.UpdateAdminTenantExtractConfigurationRequest{ + Enabled: opts.Enabled, APIBase: opts.ProviderAPIBase, Model: opts.ProviderModel, Protocol: protocol, Prompt: opts.Prompt, + } + if providerChange { + key := opts.ProviderAPIKey + request.APIKey = &key + } + return request, nil +} + +func validateEmbeddingUpdate(opts UpdateEmbeddingOptions) (apifs.UpdateAdminTenantEmbeddingConfigurationRequest, error) { + request := apifs.UpdateAdminTenantEmbeddingConfigurationRequest{Enabled: opts.Enabled} + if !opts.Enabled { + if opts.ProviderAPIBase != nil || opts.ProviderModel != nil { + return request, usageError("fs.invalid_embedding_configuration", "--enabled false cannot be combined with provider options") + } + return request, nil + } + if opts.ProviderAPIBase == nil || opts.ProviderModel == nil || !opts.ProviderAPIKeySupplied { + return request, providerTrioError() + } + apiBase := strings.TrimSpace(*opts.ProviderAPIBase) + model := strings.TrimSpace(*opts.ProviderModel) + apiKey := strings.TrimSpace(opts.ProviderAPIKey) + if err := validateProviderFields(apiBase, model, apiKey); err != nil { + return request, err + } + request.APIBase = &apiBase + request.Model = &model + request.APIKey = &apiKey + return request, nil +} + +func validateProviderFields(apiBase, model, apiKey string) error { + apiBase = strings.TrimSpace(apiBase) + model = strings.TrimSpace(model) + apiKey = strings.TrimSpace(apiKey) + if apiBase == "" || model == "" || apiKey == "" { + return providerTrioError() + } + if strings.Contains(apiKey, "*") { + return usageError("fs.invalid_ai_provider_key", "TI_FS_AI_PROVIDER_API_KEY must contain a plaintext provider key, not a masked value") + } + parsed, err := url.Parse(apiBase) + if err != nil || parsed.Scheme != "https" || parsed.Host == "" || parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" { + return usageError("fs.invalid_provider_api_base", "--provider-api-base must be a valid HTTPS URL without user information, query parameters, or a fragment") + } + return nil +} + +func validateResponse(source string, maskedAPIKey *string) error { + if strings.TrimSpace(source) == "" { + return apperr.New("fs.api_contract", "api", 1, "file system AI configuration response did not include source") + } + if maskedAPIKey != nil && *maskedAPIKey != "" && !strings.Contains(*maskedAPIKey, "*") { + return apperr.New("fs.api_contract", "api", 1, "file system AI configuration response returned an unmasked provider API key") + } + return nil +} + +func validateFileSystemID(value string) (string, error) { + if strings.TrimSpace(value) == "" { + return "", usageError("fs.missing_file_system_id", "--file-system-id is required; AI configuration requires TiDB Cloud API credentials") + } + return fscred.ValidateFileSystemID(value) +} + +func validateMediaType(value string) (string, error) { + value = strings.ToLower(strings.TrimSpace(value)) + switch value { + case MediaTypeImage, MediaTypeAudio, MediaTypeVideo: + return value, nil + default: + return "", usageError("fs.invalid_media_type", "--media-type must be one of image, audio, or video") + } +} + +func providerTrioError() error { + return usageError("fs.incomplete_ai_provider", "--provider-api-base, --provider-model, and TI_FS_AI_PROVIDER_API_KEY must be provided together") +} + +func usageError(code, message string) error { + return apperr.New(code, "usage", 2, message) +} + +func mapError(err error, operation, kind, fileSystemID, mediaType string) error { + var apiErr *api.Error + if !errors.As(err, &apiErr) { + return err + } + if apiErr.StatusCode == http.StatusNotFound { + return apperr.Wrap("fs.remote_file_system_not_found", "api", 5, fmt.Sprintf("file system %q was not found in the selected region", fileSystemID), err) + } + if kind == "embedding" && apiErr.StatusCode == http.StatusConflict { + return apperr.Wrap("fs.embedding_configuration_not_applicable", "api", 1, "embedding configuration is not applicable because this file system uses database-managed auto embedding (source=database_auto)", err) + } + if operation == "update" { + lower := strings.ToLower(apiErr.Message) + if apiErr.StatusCode == http.StatusBadRequest && strings.Contains(lower, "provider") { + return apperr.Wrap("fs.ai_provider_validation_failed", "api", 1, apiErr.Message, err) + } + if apiErr.Code == "api.network_error" || apiErr.StatusCode == http.StatusTooManyRequests || apiErr.StatusCode == http.StatusBadGateway || apiErr.StatusCode == http.StatusGatewayTimeout { + describe := "ti fs describe-file-system-embedding-configuration --file-system-id " + fileSystemID + if kind == "extract" { + describe = "ti fs describe-file-system-extract-configuration --file-system-id " + fileSystemID + " --media-type " + mediaType + } + return apperr.Wrap("fs.ai_configuration_update_ambiguous", "api", 1, "the update outcome is unknown and provider validation might have incurred a charge; run `"+describe+"` before deciding whether to retry", err) + } + } + return err +} + +func extractResult(fileSystemID, mediaType string, response apifs.AdminTenantExtractConfiguration) ExtractResult { + return ExtractResult{ + FileSystemID: fileSystemID, MediaType: mediaType, Enabled: response.Enabled, APIBase: response.APIBase, + APIKey: response.APIKey, Model: response.Model, Protocol: response.Protocol, Prompt: response.Prompt, + Source: response.Source, UpdatedAt: response.UpdatedAt, + } +} + +func embeddingResult(fileSystemID string, response apifs.AdminTenantEmbeddingConfiguration) EmbeddingResult { + return EmbeddingResult{ + FileSystemID: fileSystemID, Enabled: response.Enabled, APIBase: response.APIBase, APIKey: response.APIKey, + Model: response.Model, Source: response.Source, Generation: response.Generation, UpdatedAt: response.UpdatedAt, + } +} + +func extractDryRunBody(request apifs.UpdateAdminTenantExtractConfigurationRequest, keySupplied bool) map[string]any { + body := map[string]any{"provider_api_key_supplied": keySupplied} + addPointer(body, "enabled", request.Enabled) + addPointer(body, "api_base", request.APIBase) + addPointer(body, "model", request.Model) + addPointer(body, "protocol", request.Protocol) + addPointer(body, "prompt", request.Prompt) + if keySupplied { + body["api_key"] = "[REDACTED]" + } + return body +} + +func embeddingDryRunBody(request apifs.UpdateAdminTenantEmbeddingConfigurationRequest, keySupplied bool) map[string]any { + body := map[string]any{"enabled": request.Enabled, "provider_api_key_supplied": keySupplied} + addPointer(body, "api_base", request.APIBase) + addPointer(body, "model", request.Model) + if keySupplied { + body["api_key"] = "[REDACTED]" + } + return body +} + +func addPointer[T any](body map[string]any, name string, value *T) { + if value != nil { + body[name] = *value + } +} diff --git a/internal/fs/aiconfig/service_test.go b/internal/fs/aiconfig/service_test.go new file mode 100644 index 0000000..0c0ec79 --- /dev/null +++ b/internal/fs/aiconfig/service_test.go @@ -0,0 +1,220 @@ +package aiconfig + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/tidbcloud/ti-cli/internal/api/endpoints" + "github.com/tidbcloud/ti-cli/internal/apperr" + "github.com/tidbcloud/ti-cli/internal/config" +) + +func TestExtractConfigurationValidation(t *testing.T) { + falseValue := false + trueValue := true + base := "https://provider.example/v1" + model := "model" + qwen := ProtocolQwenASR + tooLargePrompt := strings.Repeat("x", maxPromptBytes+1) + tests := []struct { + name string + opts UpdateExtractOptions + code string + }{ + {name: "no changes", opts: UpdateExtractOptions{MediaType: MediaTypeImage}, code: "fs.extract_configuration_no_changes"}, + {name: "text media", opts: UpdateExtractOptions{MediaType: "text", Enabled: &falseValue}, code: "fs.invalid_media_type"}, + {name: "disable with prompt", opts: UpdateExtractOptions{MediaType: MediaTypeImage, Enabled: &falseValue, Prompt: stringPointer("clear")}, code: "fs.invalid_extract_configuration"}, + {name: "partial trio", opts: UpdateExtractOptions{MediaType: MediaTypeImage, Enabled: &trueValue, ProviderAPIBase: &base}, code: "fs.incomplete_ai_provider"}, + {name: "non https", opts: UpdateExtractOptions{MediaType: MediaTypeImage, ProviderAPIBase: stringPointer("http://provider.example"), ProviderModel: &model, ProviderAPIKey: "secret", ProviderAPIKeySupplied: true}, code: "fs.invalid_provider_api_base"}, + {name: "qwen image", opts: UpdateExtractOptions{MediaType: MediaTypeImage, ProviderProtocol: &qwen}, code: "fs.invalid_provider_protocol"}, + {name: "qwen audio needs trio", opts: UpdateExtractOptions{MediaType: MediaTypeAudio, ProviderProtocol: &qwen}, code: "fs.incomplete_ai_provider"}, + {name: "prompt too large", opts: UpdateExtractOptions{MediaType: MediaTypeVideo, Prompt: &tooLargePrompt}, code: "fs.invalid_extract_prompt"}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + _, err := validateExtractUpdate(test.opts) + if err == nil || apperr.CodeFor(err) != test.code { + t.Fatalf("error = %v (%s), want %s", err, apperr.CodeFor(err), test.code) + } + }) + } + + request, err := validateExtractUpdate(UpdateExtractOptions{ + MediaType: MediaTypeAudio, Enabled: &trueValue, ProviderAPIBase: stringPointer(" " + base + " "), ProviderModel: stringPointer(" " + model + " "), + ProviderAPIKey: " secret ", ProviderAPIKeySupplied: true, + }) + if err != nil { + t.Fatal(err) + } + if request.Protocol == nil || *request.Protocol != ProtocolOpenAI || request.APIBase == nil || *request.APIBase != base || request.Model == nil || *request.Model != model || request.APIKey == nil || *request.APIKey != "secret" { + t.Fatalf("request = %#v", request) + } + prompt := "" + request, err = validateExtractUpdate(UpdateExtractOptions{MediaType: MediaTypeImage, Prompt: &prompt}) + if err != nil || request.Prompt == nil || *request.Prompt != "" { + t.Fatalf("empty prompt request = %#v, %v", request, err) + } +} + +func TestEmbeddingConfigurationValidation(t *testing.T) { + base := "https://provider.example/v1" + model := "embedding-model" + if _, err := validateEmbeddingUpdate(UpdateEmbeddingOptions{Enabled: true, ProviderAPIBase: &base, ProviderModel: &model}); apperr.CodeFor(err) != "fs.incomplete_ai_provider" { + t.Fatalf("missing key error = %v", err) + } + if _, err := validateEmbeddingUpdate(UpdateEmbeddingOptions{Enabled: false, ProviderModel: &model}); apperr.CodeFor(err) != "fs.invalid_embedding_configuration" { + t.Fatalf("disable provider error = %v", err) + } + request, err := validateEmbeddingUpdate(UpdateEmbeddingOptions{Enabled: true, ProviderAPIBase: &base, ProviderModel: &model, ProviderAPIKey: "secret", ProviderAPIKeySupplied: true}) + if err != nil || request.APIKey == nil || *request.APIKey != "secret" { + t.Fatalf("request = %#v, %v", request, err) + } +} + +func TestAIConfigurationServiceOutputDryRunAndRedaction(t *testing.T) { + const secret = "provider-secret" + requests := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requests++ + assertCredentialHeaders(t, r) + switch { + case r.Method == http.MethodGet && r.URL.Path == "/v1/admin/tenants/tenant-1/extract-config/image": + _ = json.NewEncoder(w).Encode(map[string]any{"enabled": false, "source": "none"}) + case r.Method == http.MethodPut && r.URL.Path == "/v1/admin/tenants/tenant-1/extract-config/image": + var body map[string]any + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Fatal(err) + } + if body["api_key"] != secret { + t.Fatalf("body = %#v", body) + } + base, key, model, protocol := "https://provider.example/v1", "pro********", "vision-model", ProtocolOpenAI + _ = json.NewEncoder(w).Encode(map[string]any{"enabled": true, "api_base": base, "api_key": key, "model": model, "protocol": protocol, "source": "custom"}) + default: + t.Fatalf("unexpected request %s %s", r.Method, r.URL.Path) + } + })) + defer server.Close() + + service := testService(server.URL) + profile := testProfile() + described, err := service.DescribeExtract(context.Background(), DescribeExtractOptions{Profile: profile, FileSystemID: "tenant-1", MediaType: MediaTypeImage}) + if err != nil || described.Source != "none" || !strings.Contains(described.Human(), "Provider API key: none") { + t.Fatalf("describe = %#v, %v, text=%q", described, err, described.Human()) + } + enabled := true + base, model := "https://provider.example/v1", "vision-model" + opts := UpdateExtractOptions{ + Profile: profile, FileSystemID: "tenant-1", MediaType: MediaTypeImage, Enabled: &enabled, + ProviderAPIBase: &base, ProviderModel: &model, ProviderAPIKey: secret, ProviderAPIKeySupplied: true, + } + updated, err := service.UpdateExtract(context.Background(), opts) + if err != nil || updated.APIKey == nil || *updated.APIKey != "pro********" { + t.Fatalf("update = %#v, %v", updated, err) + } + beforeDryRun := requests + dryRun, err := service.DryRunUpdateExtract("ti fs update-file-system-extract-configuration", opts) + if err != nil { + t.Fatal(err) + } + encoded, err := json.Marshal(dryRun) + if err != nil { + t.Fatal(err) + } + if requests != beforeDryRun || strings.Contains(string(encoded), secret) || !strings.Contains(string(encoded), "[REDACTED]") { + t.Fatalf("dry-run requests=%d body=%s", requests-beforeDryRun, encoded) + } + for _, check := range []string{"config_and_credentials", "endpoint_selection", "permission_requirement", "provider_validation", "remote_mutation"} { + if !strings.Contains(string(encoded), check) { + t.Fatalf("dry-run body does not contain %q: %s", check, encoded) + } + } +} + +func TestAIConfigurationErrorMappings(t *testing.T) { + t.Run("database auto conflict", func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + http.Error(w, `{"error":"database auto embedding"}`, http.StatusConflict) + })) + defer server.Close() + _, err := testService(server.URL).UpdateEmbedding(context.Background(), UpdateEmbeddingOptions{Profile: testProfile(), FileSystemID: "tenant-1", Enabled: false}) + if apperr.CodeFor(err) != "fs.embedding_configuration_not_applicable" { + t.Fatalf("error = %v (%s)", err, apperr.CodeFor(err)) + } + }) + + t.Run("provider secret is redacted", func(t *testing.T) { + const secret = "provider-secret" + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + _, _ = w.Write([]byte(`{"error":"provider validation failed for provider-secret"}`)) + })) + defer server.Close() + enabled := true + base, model := "https://provider.example/v1", "model" + _, err := testService(server.URL).UpdateEmbedding(context.Background(), UpdateEmbeddingOptions{ + Profile: testProfile(), FileSystemID: "tenant-1", Enabled: enabled, ProviderAPIBase: &base, + ProviderModel: &model, ProviderAPIKey: secret, ProviderAPIKeySupplied: true, + }) + if err == nil || apperr.CodeFor(err) != "fs.ai_provider_validation_failed" || strings.Contains(err.Error(), secret) { + t.Fatalf("error = %v (%s)", err, apperr.CodeFor(err)) + } + }) + + for _, test := range []struct { + name string + body string + }{ + {name: "missing source", body: `{"enabled":false}`}, + {name: "unmasked key", body: `{"enabled":true,"source":"custom","api_key":"provider-secret"}`}, + } { + t.Run(test.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = w.Write([]byte(test.body)) + })) + defer server.Close() + _, err := testService(server.URL).DescribeEmbedding(context.Background(), DescribeEmbeddingOptions{Profile: testProfile(), FileSystemID: "tenant-1"}) + if apperr.CodeFor(err) != "fs.api_contract" { + t.Fatalf("error = %v (%s)", err, apperr.CodeFor(err)) + } + }) + } +} + +func TestAIConfigurationResponseSources(t *testing.T) { + masked := "key********" + for _, source := range []string{"custom", "default", "none", "database_auto", "future_source"} { + if err := validateResponse(source, &masked); err != nil { + t.Fatalf("source %q rejected: %v", source, err) + } + } +} + +func testService(baseURL string) Service { + return Service{Resolver: endpoints.Resolver{FSManifest: &endpoints.FSRegionManifest{Regions: []endpoints.FSRegionManifestEntry{{ + RegionCode: "aws-us-east-1", Mode: endpoints.DefaultFSMode, ServerURL: baseURL, CloudProvider: "aws", TiDBRegion: "us-east-1", + }}}}} +} + +func testProfile() *config.Profile { + return &config.Profile{Name: "test", PlacementRegionCode: "aws-us-east-1", CloudProvider: "aws", RegionCode: "us-east-1", TiDBCloudPublicKey: "public", TiDBCloudPrivateKey: "private"} +} + +func assertCredentialHeaders(t *testing.T, request *http.Request) { + t.Helper() + if request.Header.Get("X-TiDBCloud-Public-Key") != "public" || request.Header.Get("X-TiDBCloud-Private-Key") != "private" { + t.Fatalf("credential headers = %q/%q", request.Header.Get("X-TiDBCloud-Public-Key"), request.Header.Get("X-TiDBCloud-Private-Key")) + } + if request.Header.Get("Authorization") != "" { + t.Fatalf("unexpected bearer header %q", request.Header.Get("Authorization")) + } +} + +func stringPointer(value string) *string { + return &value +} diff --git a/internal/fs/control.go b/internal/fs/control.go index 572a09a..81f47ac 100644 --- a/internal/fs/control.go +++ b/internal/fs/control.go @@ -381,7 +381,23 @@ func (r DescribeFileSystemResult) Human() string { fmt.Sprintf("Local token: %t", r.HasLocalToken), } if r.Quota != nil { - lines = append(lines, fmt.Sprintf("Quota: %v", r.Quota)) + spendingLimit := "none" + if r.Quota.Config.TiDBCloudSpendingLimit != nil { + spendingLimit = fmt.Sprintf("%d", *r.Quota.Config.TiDBCloudSpendingLimit) + } + lines = append(lines, + fmt.Sprintf("Quota max storage bytes: %d", r.Quota.Config.MaxStorageSize), + fmt.Sprintf("Quota max file size bytes: %d", r.Quota.Config.MaxFileSize), + fmt.Sprintf("Quota max file count: %d", r.Quota.Config.MaxFileCount), + fmt.Sprintf("Quota max media LLM files: %d", r.Quota.Config.MaxMediaLLMFiles), + fmt.Sprintf("Quota max video LLM files: %d", r.Quota.Config.MaxVideoLLMFiles), + "Quota TiDB Cloud spending limit: "+spendingLimit, + fmt.Sprintf("Usage storage bytes: %d", r.Quota.Usage.StorageBytes), + fmt.Sprintf("Usage reserved bytes: %d", r.Quota.Usage.ReservedBytes), + fmt.Sprintf("Usage file count: %d", r.Quota.Usage.FileCount), + fmt.Sprintf("Usage media file count: %d", r.Quota.Usage.MediaFileCount), + fmt.Sprintf("Usage video file count: %d", r.Quota.Usage.VideoFileCount), + ) } return strings.Join(lines, "\n") } diff --git a/internal/fs/tenant_control_test.go b/internal/fs/tenant_control_test.go index 9eb799b..be46015 100644 --- a/internal/fs/tenant_control_test.go +++ b/internal/fs/tenant_control_test.go @@ -209,6 +209,12 @@ func TestTenantControlDescribeAndDeleteUseIDs(t *testing.T) { if described.DisplayName != "workspace-one" || described.Labels["team"] != "ai" || !described.HasLocalToken || !strings.Contains(described.Human(), "Labels: team=ai") { t.Fatalf("describe = %#v", described) } + if described.Quota == nil || described.Quota.Config.MaxMediaLLMFiles != 100 || described.Quota.Config.MaxVideoLLMFiles != 30 || described.Quota.Usage.MediaFileCount != 12 || described.Quota.Usage.VideoFileCount != 3 { + t.Fatalf("describe quota = %#v", described.Quota) + } + if text := described.Human(); !strings.Contains(text, "Quota max media LLM files: 100") || !strings.Contains(text, "Usage video file count: 3") { + t.Fatalf("describe text = %q", text) + } deleted, err := service.DeleteFileSystem(context.Background(), DeleteFileSystemOptions{Profile: profile, FileSystemID: "tenant-1"}) if err != nil { t.Fatal(err) @@ -389,8 +395,8 @@ func assertAdminCredentialHeaders(t *testing.T, request *http.Request) { func tenantQuotaFixture() map[string]any { return map[string]any{ - "config": map[string]any{"max_storage_size": 1024, "max_file_size": 128, "max_file_count": 1000, "tidbcloud_spending_limit": nil}, - "usage": map[string]any{"storage_bytes": 12, "reserved_bytes": 3, "file_count": 2}, + "config": map[string]any{"max_storage_size": 1024, "max_file_size": 128, "max_file_count": 1000, "max_media_llm_files": 100, "max_video_llm_files": 30, "tidbcloud_spending_limit": nil}, + "usage": map[string]any{"storage_bytes": 12, "reserved_bytes": 3, "file_count": 2, "media_file_count": 12, "video_file_count": 3}, } }