mirror of
https://github.com/coder/coder.git
synced 2026-06-02 20:48:20 +00:00
d8ff67fb68
## Summary
Adds the database schema, API endpoints, SDK types, and encryption
wrappers for admin-managed MCP (Model Context Protocol) server
configurations that chatd can consume. This is the backend foundation
for allowing external MCP tools (Sentry, Linear, GitHub, etc.) to be
used during AI chat sessions.
## Database
Two new tables:
- **`mcp_server_configs`**: Admin-managed server definitions with URL,
transport (Streamable HTTP / SSE), auth config (none / OAuth2 / API key
/ custom headers), tool allow/deny lists, and an availability policy
(`force_on` / `default_on` / `default_off`). Includes CHECK constraints
on transport, auth_type, and availability values.
- **`mcp_server_user_tokens`**: Per-user OAuth2 tokens for servers
requiring individual authentication. Cascades on user/config deletion.
New column on `chats` table:
- **`mcp_server_ids UUID[]`**: Per-chat MCP server selection, following
the same pattern as `model_config_id` — passed at chat creation,
changeable per-message with nil-means-no-change semantics.
## API Endpoints
All routes are under `/api/experimental/mcp/servers/` and gated behind
the `agents` experiment.
**Admin endpoints** (`ResourceDeploymentConfig` auth):
- `POST /` — Create MCP server config
- `PATCH /{id}` — Update MCP server config (full-replace)
- `DELETE /{id}` — Delete MCP server config
**Authenticated endpoints** (all users, enabled servers only for
non-admins):
- `GET /` — List configs (admins see all, members see enabled-only with
admin fields redacted)
- `GET /{id}` — Get config by ID (with `auth_connected` populated
per-user)
**OAuth2 per-user auth flow:**
- `GET /{id}/oauth2/connect` — Initiate OAuth2 flow (state cookie CSRF
protection)
- `GET /{id}/oauth2/callback` — Handle OAuth2 callback, store tokens
- `DELETE /{id}/oauth2/disconnect` — Remove stored OAuth2 tokens
## Security
- **Secrets never returned**: `OAuth2ClientSecret`, `APIKeyValue`, and
`CustomHeaders` are never in API responses — only boolean indicators
(`has_oauth2_secret`, `has_api_key`, `has_custom_headers`).
- **Field redaction for non-admins**: `convertMCPServerConfigRedacted`
strips `OAuth2ClientID`, auth URLs, scopes, and `APIKeyHeader` from
non-admin responses.
- **dbcrypt encryption at rest**: All 5 secret fields use `dbcrypt_keys`
encryption with full encrypt-on-write / decrypt-on-read wrappers (11
dbcrypt method overrides + 2 helpers), following the same pattern as
`chat_providers.api_key`.
- **OAuth2 CSRF protection**: State parameter stored in `HttpOnly`
cookie with `HTTPCookies.Apply()` for correct `Secure`/`SameSite` behind
TLS-terminating proxies.
- **dbauthz authorization**: All 18 querier methods have authorization
wrappers. Read operations use `ActionRead`, write operations use
`ActionUpdate` on `ResourceDeploymentConfig`.
## Governance Model
| Control | Implementation |
|---------|---------------|
| **Global kill switch** | `enabled` defaults to `false` |
| **Availability policy** | `force_on` (always injected), `default_on`
(pre-selected), `default_off` (opt-in) |
| **Per-chat selection** | `mcp_server_ids` on `CreateChatRequest` /
`CreateChatMessageRequest` |
| **Auth gate** | OAuth2 servers require per-user auth before tools are
injected |
| **Tool-level allow/deny** | Arrays on `mcp_server_configs` for
granular tool filtering |
| **Secrets encrypted at rest** | Uses `dbcrypt_keys` (same pattern as
`chat_providers.api_key`) |
## Tests
8 test functions covering:
- Full CRUD lifecycle (create, list, update, delete)
- Non-admin visibility filtering (enabled-only, field redaction)
- `auth_connected` population for OAuth2 vs non-OAuth2 servers
- Availability policy validation (valid values + invalid rejection)
- Unique slug enforcement (409 Conflict)
- OAuth2 disconnect idempotency
- Chat creation with `mcp_server_ids` persistence
## Known Limitations (Deferred)
These are documented and intentional for an experimental feature:
- **Audit logging** not yet wired — will add when feature stabilizes
- **Cross-field validation** (e.g., OAuth2 fields required when
`auth_type=oauth2`) — admin-only endpoint, will add when stabilizing
- **`force_on` auto-injection** — query exists but not yet wired into
chatd tool injection (follow-up)
- **Additional test coverage** — 403 auth tests, GET-by-ID tests,
callback CSRF tests planned for follow-up
## What's NOT in this PR
- Frontend UI (admin panel + chat picker)
- Actual MCP client connections (`chatd/chatmcp/` manager)
- Tool injection into `chatloop/`
882 lines
24 KiB
SQL
882 lines
24 KiB
SQL
-- name: ArchiveChatByID :exec
|
|
UPDATE chats SET archived = true, updated_at = NOW()
|
|
WHERE id = @id OR root_chat_id = @id;
|
|
|
|
-- name: UnarchiveChatByID :exec
|
|
UPDATE chats SET archived = false, updated_at = NOW() WHERE id = @id::uuid;
|
|
|
|
-- name: SoftDeleteChatMessagesAfterID :exec
|
|
UPDATE
|
|
chat_messages
|
|
SET
|
|
deleted = true
|
|
WHERE
|
|
chat_id = @chat_id::uuid
|
|
AND id > @after_id::bigint;
|
|
|
|
-- name: SoftDeleteChatMessageByID :exec
|
|
UPDATE
|
|
chat_messages
|
|
SET
|
|
deleted = true
|
|
WHERE
|
|
id = @id::bigint;
|
|
|
|
-- name: GetChatByID :one
|
|
SELECT
|
|
*
|
|
FROM
|
|
chats
|
|
WHERE
|
|
id = @id::uuid;
|
|
|
|
-- name: GetChatMessageByID :one
|
|
SELECT
|
|
*
|
|
FROM
|
|
chat_messages
|
|
WHERE
|
|
id = @id::bigint
|
|
AND deleted = false;
|
|
|
|
-- name: GetChatMessagesByChatID :many
|
|
SELECT
|
|
*
|
|
FROM
|
|
chat_messages
|
|
WHERE
|
|
chat_id = @chat_id::uuid
|
|
AND id > @after_id::bigint
|
|
AND visibility IN ('user', 'both')
|
|
AND deleted = false
|
|
ORDER BY
|
|
created_at ASC;
|
|
|
|
-- name: GetChatMessagesByChatIDDescPaginated :many
|
|
SELECT
|
|
*
|
|
FROM
|
|
chat_messages
|
|
WHERE
|
|
chat_id = @chat_id::uuid
|
|
AND CASE
|
|
WHEN @before_id::bigint > 0 THEN id < @before_id::bigint
|
|
ELSE true
|
|
END
|
|
AND visibility IN ('user', 'both')
|
|
AND deleted = false
|
|
ORDER BY
|
|
id DESC
|
|
LIMIT
|
|
COALESCE(NULLIF(@limit_val::int, 0), 50);
|
|
|
|
-- name: GetChatMessagesForPromptByChatID :many
|
|
WITH latest_compressed_summary AS (
|
|
SELECT
|
|
id
|
|
FROM
|
|
chat_messages
|
|
WHERE
|
|
chat_id = @chat_id::uuid
|
|
AND compressed = TRUE
|
|
AND deleted = false
|
|
AND visibility = 'model'
|
|
ORDER BY
|
|
created_at DESC,
|
|
id DESC
|
|
LIMIT
|
|
1
|
|
)
|
|
SELECT
|
|
*
|
|
FROM
|
|
chat_messages
|
|
WHERE
|
|
chat_id = @chat_id::uuid
|
|
AND visibility IN ('model', 'both')
|
|
AND deleted = false
|
|
AND (
|
|
(
|
|
role = 'system'
|
|
AND compressed = FALSE
|
|
)
|
|
OR (
|
|
compressed = FALSE
|
|
AND (
|
|
NOT EXISTS (
|
|
SELECT
|
|
1
|
|
FROM
|
|
latest_compressed_summary
|
|
)
|
|
OR id > (
|
|
SELECT
|
|
id
|
|
FROM
|
|
latest_compressed_summary
|
|
)
|
|
)
|
|
)
|
|
OR id = (
|
|
SELECT
|
|
id
|
|
FROM
|
|
latest_compressed_summary
|
|
)
|
|
)
|
|
ORDER BY
|
|
created_at ASC,
|
|
id ASC;
|
|
|
|
-- name: GetChats :many
|
|
SELECT
|
|
*
|
|
FROM
|
|
chats
|
|
WHERE
|
|
CASE
|
|
WHEN @owner_id :: uuid != '00000000-0000-0000-0000-000000000000'::uuid THEN chats.owner_id = @owner_id
|
|
ELSE true
|
|
END
|
|
AND CASE
|
|
WHEN sqlc.narg('archived') :: boolean IS NULL THEN true
|
|
ELSE chats.archived = sqlc.narg('archived') :: boolean
|
|
END
|
|
AND CASE
|
|
-- This allows using the last element on a page as effectively a cursor.
|
|
-- This is an important option for scripts that need to paginate without
|
|
-- duplicating or missing data.
|
|
WHEN @after_id :: uuid != '00000000-0000-0000-0000-000000000000'::uuid THEN (
|
|
-- The pagination cursor is the last ID of the previous page.
|
|
-- The query is ordered by the updated_at field, so select all
|
|
-- rows before the cursor.
|
|
(updated_at, id) < (
|
|
SELECT
|
|
updated_at, id
|
|
FROM
|
|
chats
|
|
WHERE
|
|
id = @after_id
|
|
)
|
|
)
|
|
ELSE true
|
|
END
|
|
-- Authorize Filter clause will be injected below in GetAuthorizedChats
|
|
-- @authorize_filter
|
|
ORDER BY
|
|
-- Deterministic and consistent ordering of all rows, even if they share
|
|
-- a timestamp. This is to ensure consistent pagination.
|
|
(updated_at, id) DESC OFFSET @offset_opt
|
|
LIMIT
|
|
-- The chat list is unbounded and expected to grow large.
|
|
-- Default to 50 to prevent accidental excessively large queries.
|
|
COALESCE(NULLIF(@limit_opt :: int, 0), 50);
|
|
|
|
-- name: InsertChat :one
|
|
INSERT INTO chats (
|
|
owner_id,
|
|
workspace_id,
|
|
parent_chat_id,
|
|
root_chat_id,
|
|
last_model_config_id,
|
|
title,
|
|
mode,
|
|
mcp_server_ids
|
|
) VALUES (
|
|
@owner_id::uuid,
|
|
sqlc.narg('workspace_id')::uuid,
|
|
sqlc.narg('parent_chat_id')::uuid,
|
|
sqlc.narg('root_chat_id')::uuid,
|
|
@last_model_config_id::uuid,
|
|
@title::text,
|
|
sqlc.narg('mode')::chat_mode,
|
|
COALESCE(@mcp_server_ids::uuid[], '{}'::uuid[])
|
|
)
|
|
RETURNING
|
|
*;
|
|
|
|
-- name: InsertChatMessages :many
|
|
WITH updated_chat AS (
|
|
UPDATE
|
|
chats
|
|
SET
|
|
last_model_config_id = (
|
|
SELECT val
|
|
FROM UNNEST(@model_config_id::uuid[])
|
|
WITH ORDINALITY AS t(val, ord)
|
|
WHERE val != '00000000-0000-0000-0000-000000000000'::uuid
|
|
ORDER BY ord DESC
|
|
LIMIT 1
|
|
)
|
|
WHERE
|
|
id = @chat_id::uuid
|
|
AND EXISTS (
|
|
SELECT 1
|
|
FROM UNNEST(@model_config_id::uuid[])
|
|
WHERE unnest != '00000000-0000-0000-0000-000000000000'::uuid
|
|
)
|
|
AND chats.last_model_config_id IS DISTINCT FROM (
|
|
SELECT val
|
|
FROM UNNEST(@model_config_id::uuid[])
|
|
WITH ORDINALITY AS t(val, ord)
|
|
WHERE val != '00000000-0000-0000-0000-000000000000'::uuid
|
|
ORDER BY ord DESC
|
|
LIMIT 1
|
|
)
|
|
)
|
|
INSERT INTO chat_messages (
|
|
chat_id,
|
|
created_by,
|
|
model_config_id,
|
|
role,
|
|
content,
|
|
content_version,
|
|
visibility,
|
|
input_tokens,
|
|
output_tokens,
|
|
total_tokens,
|
|
reasoning_tokens,
|
|
cache_creation_tokens,
|
|
cache_read_tokens,
|
|
context_limit,
|
|
compressed,
|
|
total_cost_micros,
|
|
runtime_ms
|
|
)
|
|
SELECT
|
|
@chat_id::uuid,
|
|
NULLIF(UNNEST(@created_by::uuid[]), '00000000-0000-0000-0000-000000000000'::uuid),
|
|
NULLIF(UNNEST(@model_config_id::uuid[]), '00000000-0000-0000-0000-000000000000'::uuid),
|
|
UNNEST(@role::chat_message_role[]),
|
|
UNNEST(@content::text[])::jsonb,
|
|
UNNEST(@content_version::smallint[]),
|
|
UNNEST(@visibility::chat_message_visibility[]),
|
|
NULLIF(UNNEST(@input_tokens::bigint[]), 0),
|
|
NULLIF(UNNEST(@output_tokens::bigint[]), 0),
|
|
NULLIF(UNNEST(@total_tokens::bigint[]), 0),
|
|
NULLIF(UNNEST(@reasoning_tokens::bigint[]), 0),
|
|
NULLIF(UNNEST(@cache_creation_tokens::bigint[]), 0),
|
|
NULLIF(UNNEST(@cache_read_tokens::bigint[]), 0),
|
|
NULLIF(UNNEST(@context_limit::bigint[]), 0),
|
|
UNNEST(@compressed::boolean[]),
|
|
NULLIF(UNNEST(@total_cost_micros::bigint[]), 0),
|
|
NULLIF(UNNEST(@runtime_ms::bigint[]), 0)
|
|
RETURNING
|
|
*;
|
|
|
|
-- name: UpdateChatMessageByID :one
|
|
UPDATE
|
|
chat_messages
|
|
SET
|
|
model_config_id = COALESCE(sqlc.narg('model_config_id')::uuid, model_config_id),
|
|
content = sqlc.narg('content')::jsonb
|
|
WHERE
|
|
id = @id::bigint
|
|
RETURNING
|
|
*;
|
|
|
|
-- name: UpdateChatByID :one
|
|
UPDATE
|
|
chats
|
|
SET
|
|
title = @title::text,
|
|
updated_at = NOW()
|
|
WHERE
|
|
id = @id::uuid
|
|
RETURNING
|
|
*;
|
|
|
|
-- name: UpdateChatWorkspace :one
|
|
UPDATE
|
|
chats
|
|
SET
|
|
workspace_id = sqlc.narg('workspace_id')::uuid,
|
|
updated_at = NOW()
|
|
WHERE
|
|
id = @id::uuid
|
|
RETURNING
|
|
*;
|
|
|
|
-- name: UpdateChatMCPServerIDs :one
|
|
UPDATE
|
|
chats
|
|
SET
|
|
mcp_server_ids = @mcp_server_ids::uuid[],
|
|
updated_at = NOW()
|
|
WHERE
|
|
id = @id::uuid
|
|
RETURNING
|
|
*;
|
|
|
|
-- name: AcquireChats :many
|
|
-- Acquires up to @num_chats pending chats for processing. Uses SKIP LOCKED
|
|
-- to prevent multiple replicas from acquiring the same chat.
|
|
UPDATE
|
|
chats
|
|
SET
|
|
status = 'running'::chat_status,
|
|
started_at = @started_at::timestamptz,
|
|
heartbeat_at = @started_at::timestamptz,
|
|
updated_at = @started_at::timestamptz,
|
|
worker_id = @worker_id::uuid
|
|
WHERE
|
|
id = ANY(
|
|
SELECT
|
|
id
|
|
FROM
|
|
chats
|
|
WHERE
|
|
status = 'pending'::chat_status
|
|
ORDER BY
|
|
updated_at ASC
|
|
FOR UPDATE
|
|
SKIP LOCKED
|
|
LIMIT
|
|
@num_chats::int
|
|
)
|
|
RETURNING
|
|
*;
|
|
|
|
-- name: UpdateChatStatus :one
|
|
UPDATE
|
|
chats
|
|
SET
|
|
status = @status::chat_status,
|
|
worker_id = sqlc.narg('worker_id')::uuid,
|
|
started_at = sqlc.narg('started_at')::timestamptz,
|
|
heartbeat_at = sqlc.narg('heartbeat_at')::timestamptz,
|
|
last_error = sqlc.narg('last_error')::text,
|
|
updated_at = NOW()
|
|
WHERE
|
|
id = @id::uuid
|
|
RETURNING
|
|
*;
|
|
|
|
-- name: GetStaleChats :many
|
|
-- Find chats that appear stuck (running but heartbeat has expired).
|
|
-- Used for recovery after coderd crashes or long hangs.
|
|
SELECT
|
|
*
|
|
FROM
|
|
chats
|
|
WHERE
|
|
status = 'running'::chat_status
|
|
AND heartbeat_at < @stale_threshold::timestamptz;
|
|
|
|
-- name: UpdateChatHeartbeat :execrows
|
|
-- Bumps the heartbeat timestamp for a running chat so that other
|
|
-- replicas know the worker is still alive.
|
|
UPDATE
|
|
chats
|
|
SET
|
|
heartbeat_at = NOW()
|
|
WHERE
|
|
id = @id::uuid
|
|
AND worker_id = @worker_id::uuid
|
|
AND status = 'running'::chat_status;
|
|
|
|
-- name: GetChatDiffStatusByChatID :one
|
|
SELECT
|
|
*
|
|
FROM
|
|
chat_diff_statuses
|
|
WHERE
|
|
chat_id = @chat_id::uuid;
|
|
|
|
-- name: GetChatDiffStatusesByChatIDs :many
|
|
SELECT
|
|
*
|
|
FROM
|
|
chat_diff_statuses
|
|
WHERE
|
|
chat_id = ANY(@chat_ids::uuid[]);
|
|
|
|
-- name: UpsertChatDiffStatusReference :one
|
|
INSERT INTO chat_diff_statuses (
|
|
chat_id,
|
|
url,
|
|
git_branch,
|
|
git_remote_origin,
|
|
stale_at
|
|
) VALUES (
|
|
@chat_id::uuid,
|
|
sqlc.narg('url')::text,
|
|
@git_branch::text,
|
|
@git_remote_origin::text,
|
|
@stale_at::timestamptz
|
|
)
|
|
ON CONFLICT (chat_id) DO UPDATE
|
|
SET
|
|
url = CASE
|
|
WHEN EXCLUDED.url IS NOT NULL THEN EXCLUDED.url
|
|
ELSE chat_diff_statuses.url
|
|
END,
|
|
git_branch = CASE
|
|
WHEN EXCLUDED.git_branch != '' THEN EXCLUDED.git_branch
|
|
ELSE chat_diff_statuses.git_branch
|
|
END,
|
|
git_remote_origin = CASE
|
|
WHEN EXCLUDED.git_remote_origin != '' THEN EXCLUDED.git_remote_origin
|
|
ELSE chat_diff_statuses.git_remote_origin
|
|
END,
|
|
stale_at = EXCLUDED.stale_at,
|
|
updated_at = NOW()
|
|
RETURNING
|
|
*;
|
|
|
|
-- name: UpsertChatDiffStatus :one
|
|
INSERT INTO chat_diff_statuses (
|
|
chat_id,
|
|
url,
|
|
pull_request_state,
|
|
pull_request_title,
|
|
pull_request_draft,
|
|
changes_requested,
|
|
additions,
|
|
deletions,
|
|
changed_files,
|
|
author_login,
|
|
author_avatar_url,
|
|
base_branch,
|
|
head_branch,
|
|
pr_number,
|
|
commits,
|
|
approved,
|
|
reviewer_count,
|
|
refreshed_at,
|
|
stale_at
|
|
) VALUES (
|
|
@chat_id::uuid,
|
|
sqlc.narg('url')::text,
|
|
sqlc.narg('pull_request_state')::text,
|
|
@pull_request_title::text,
|
|
@pull_request_draft::boolean,
|
|
@changes_requested::boolean,
|
|
@additions::integer,
|
|
@deletions::integer,
|
|
@changed_files::integer,
|
|
sqlc.narg('author_login')::text,
|
|
sqlc.narg('author_avatar_url')::text,
|
|
sqlc.narg('base_branch')::text,
|
|
sqlc.narg('head_branch')::text,
|
|
sqlc.narg('pr_number')::integer,
|
|
sqlc.narg('commits')::integer,
|
|
sqlc.narg('approved')::boolean,
|
|
sqlc.narg('reviewer_count')::integer,
|
|
@refreshed_at::timestamptz,
|
|
@stale_at::timestamptz
|
|
)
|
|
ON CONFLICT (chat_id) DO UPDATE
|
|
SET
|
|
url = EXCLUDED.url,
|
|
pull_request_state = EXCLUDED.pull_request_state,
|
|
pull_request_title = EXCLUDED.pull_request_title,
|
|
pull_request_draft = EXCLUDED.pull_request_draft,
|
|
changes_requested = EXCLUDED.changes_requested,
|
|
additions = EXCLUDED.additions,
|
|
deletions = EXCLUDED.deletions,
|
|
changed_files = EXCLUDED.changed_files,
|
|
author_login = EXCLUDED.author_login,
|
|
author_avatar_url = EXCLUDED.author_avatar_url,
|
|
base_branch = EXCLUDED.base_branch,
|
|
head_branch = EXCLUDED.head_branch,
|
|
pr_number = EXCLUDED.pr_number,
|
|
commits = EXCLUDED.commits,
|
|
approved = EXCLUDED.approved,
|
|
reviewer_count = EXCLUDED.reviewer_count,
|
|
refreshed_at = EXCLUDED.refreshed_at,
|
|
stale_at = EXCLUDED.stale_at,
|
|
updated_at = NOW()
|
|
RETURNING
|
|
*;
|
|
|
|
-- name: InsertChatQueuedMessage :one
|
|
INSERT INTO chat_queued_messages (chat_id, content)
|
|
VALUES (@chat_id, @content)
|
|
RETURNING *;
|
|
|
|
-- name: GetChatQueuedMessages :many
|
|
SELECT * FROM chat_queued_messages
|
|
WHERE chat_id = @chat_id
|
|
ORDER BY id ASC;
|
|
|
|
-- name: DeleteChatQueuedMessage :exec
|
|
DELETE FROM chat_queued_messages WHERE id = @id AND chat_id = @chat_id;
|
|
|
|
-- name: DeleteAllChatQueuedMessages :exec
|
|
DELETE FROM chat_queued_messages WHERE chat_id = @chat_id;
|
|
|
|
-- name: PopNextQueuedMessage :one
|
|
DELETE FROM chat_queued_messages
|
|
WHERE id = (
|
|
SELECT cqm.id FROM chat_queued_messages cqm
|
|
WHERE cqm.chat_id = @chat_id
|
|
ORDER BY cqm.id ASC
|
|
LIMIT 1
|
|
)
|
|
RETURNING *;
|
|
|
|
-- name: GetLastChatMessageByRole :one
|
|
SELECT
|
|
*
|
|
FROM
|
|
chat_messages
|
|
WHERE
|
|
chat_id = @chat_id::uuid
|
|
AND role = @role::chat_message_role
|
|
AND deleted = false
|
|
ORDER BY
|
|
created_at DESC, id DESC
|
|
LIMIT
|
|
1;
|
|
|
|
-- name: GetChatByIDForUpdate :one
|
|
SELECT * FROM chats WHERE id = @id::uuid FOR UPDATE;
|
|
|
|
-- name: AcquireStaleChatDiffStatuses :many
|
|
WITH acquired AS (
|
|
UPDATE
|
|
chat_diff_statuses
|
|
SET
|
|
-- Claim for 5 minutes. The worker sets the real stale_at
|
|
-- after refresh. If the worker crashes, rows become eligible
|
|
-- again after this interval.
|
|
stale_at = NOW() + INTERVAL '5 minutes',
|
|
updated_at = NOW()
|
|
WHERE
|
|
chat_id IN (
|
|
SELECT
|
|
cds.chat_id
|
|
FROM
|
|
chat_diff_statuses cds
|
|
INNER JOIN
|
|
chats c ON c.id = cds.chat_id
|
|
WHERE
|
|
cds.stale_at <= NOW()
|
|
AND cds.git_remote_origin != ''
|
|
AND cds.git_branch != ''
|
|
AND c.archived = FALSE
|
|
ORDER BY
|
|
cds.stale_at ASC
|
|
FOR UPDATE OF cds
|
|
SKIP LOCKED
|
|
LIMIT
|
|
@limit_val::int
|
|
)
|
|
RETURNING *
|
|
)
|
|
SELECT
|
|
acquired.*,
|
|
c.owner_id
|
|
FROM
|
|
acquired
|
|
INNER JOIN
|
|
chats c ON c.id = acquired.chat_id;
|
|
|
|
-- name: BackoffChatDiffStatus :exec
|
|
UPDATE
|
|
chat_diff_statuses
|
|
SET
|
|
stale_at = @stale_at::timestamptz,
|
|
updated_at = NOW()
|
|
WHERE
|
|
chat_id = @chat_id::uuid;
|
|
|
|
-- name: GetChatCostSummary :one
|
|
-- Aggregate cost summary for a single user within a date range.
|
|
-- Only counts assistant-role messages.
|
|
SELECT
|
|
COALESCE(SUM(cm.total_cost_micros), 0)::bigint AS total_cost_micros,
|
|
COUNT(*) FILTER (
|
|
WHERE cm.total_cost_micros IS NOT NULL
|
|
)::bigint AS priced_message_count,
|
|
COUNT(*) FILTER (
|
|
WHERE cm.total_cost_micros IS NULL
|
|
AND (
|
|
cm.input_tokens IS NOT NULL
|
|
OR cm.output_tokens IS NOT NULL
|
|
OR cm.reasoning_tokens IS NOT NULL
|
|
OR cm.cache_creation_tokens IS NOT NULL
|
|
OR cm.cache_read_tokens IS NOT NULL
|
|
)
|
|
)::bigint AS unpriced_message_count,
|
|
COALESCE(SUM(cm.input_tokens), 0)::bigint AS total_input_tokens,
|
|
COALESCE(SUM(cm.output_tokens), 0)::bigint AS total_output_tokens,
|
|
COALESCE(SUM(cm.cache_read_tokens), 0)::bigint AS total_cache_read_tokens,
|
|
COALESCE(SUM(cm.cache_creation_tokens), 0)::bigint AS total_cache_creation_tokens
|
|
FROM
|
|
chat_messages cm
|
|
JOIN
|
|
chats c ON c.id = cm.chat_id
|
|
WHERE
|
|
c.owner_id = @owner_id::uuid
|
|
AND cm.role = 'assistant'
|
|
AND cm.created_at >= @start_date::timestamptz
|
|
AND cm.created_at < @end_date::timestamptz;
|
|
|
|
-- name: GetChatCostPerModel :many
|
|
-- Per-model cost breakdown for a single user within a date range.
|
|
-- Only counts assistant-role messages that have a model_config_id.
|
|
SELECT
|
|
cmc.id AS model_config_id,
|
|
cmc.display_name,
|
|
cmc.provider,
|
|
cmc.model,
|
|
COALESCE(SUM(cm.total_cost_micros), 0)::bigint AS total_cost_micros,
|
|
COUNT(*) FILTER (
|
|
WHERE cm.input_tokens IS NOT NULL
|
|
OR cm.output_tokens IS NOT NULL
|
|
OR cm.reasoning_tokens IS NOT NULL
|
|
OR cm.cache_creation_tokens IS NOT NULL
|
|
OR cm.cache_read_tokens IS NOT NULL
|
|
)::bigint AS message_count,
|
|
COALESCE(SUM(cm.input_tokens), 0)::bigint AS total_input_tokens,
|
|
COALESCE(SUM(cm.output_tokens), 0)::bigint AS total_output_tokens,
|
|
COALESCE(SUM(cm.cache_read_tokens), 0)::bigint AS total_cache_read_tokens,
|
|
COALESCE(SUM(cm.cache_creation_tokens), 0)::bigint AS total_cache_creation_tokens
|
|
FROM
|
|
chat_messages cm
|
|
JOIN
|
|
chats c ON c.id = cm.chat_id
|
|
JOIN
|
|
chat_model_configs cmc ON cmc.id = cm.model_config_id
|
|
WHERE
|
|
c.owner_id = @owner_id::uuid
|
|
AND cm.role = 'assistant'
|
|
AND cm.created_at >= @start_date::timestamptz
|
|
AND cm.created_at < @end_date::timestamptz
|
|
GROUP BY
|
|
cmc.id, cmc.display_name, cmc.provider, cmc.model
|
|
ORDER BY
|
|
total_cost_micros DESC;
|
|
|
|
-- name: GetChatCostPerChat :many
|
|
-- Per-root-chat cost breakdown for a single user within a date range.
|
|
-- Groups by root_chat_id so forked chats roll up under their root.
|
|
-- Only counts assistant-role messages.
|
|
WITH chat_costs AS (
|
|
SELECT
|
|
COALESCE(c.root_chat_id, c.id) AS root_chat_id,
|
|
COALESCE(SUM(cm.total_cost_micros), 0)::bigint AS total_cost_micros,
|
|
COUNT(*) FILTER (
|
|
WHERE cm.input_tokens IS NOT NULL
|
|
OR cm.output_tokens IS NOT NULL
|
|
OR cm.reasoning_tokens IS NOT NULL
|
|
OR cm.cache_creation_tokens IS NOT NULL
|
|
OR cm.cache_read_tokens IS NOT NULL
|
|
)::bigint AS message_count,
|
|
COALESCE(SUM(cm.input_tokens), 0)::bigint AS total_input_tokens,
|
|
COALESCE(SUM(cm.output_tokens), 0)::bigint AS total_output_tokens,
|
|
COALESCE(SUM(cm.cache_read_tokens), 0)::bigint AS total_cache_read_tokens,
|
|
COALESCE(SUM(cm.cache_creation_tokens), 0)::bigint AS total_cache_creation_tokens
|
|
FROM chat_messages cm
|
|
JOIN chats c ON c.id = cm.chat_id
|
|
WHERE c.owner_id = @owner_id::uuid
|
|
AND cm.role = 'assistant'
|
|
AND cm.created_at >= @start_date::timestamptz
|
|
AND cm.created_at < @end_date::timestamptz
|
|
GROUP BY COALESCE(c.root_chat_id, c.id)
|
|
)
|
|
SELECT
|
|
cc.root_chat_id,
|
|
COALESCE(rc.title, '') AS chat_title,
|
|
cc.total_cost_micros,
|
|
cc.message_count,
|
|
cc.total_input_tokens,
|
|
cc.total_output_tokens,
|
|
cc.total_cache_read_tokens,
|
|
cc.total_cache_creation_tokens
|
|
FROM chat_costs cc
|
|
LEFT JOIN chats rc ON rc.id = cc.root_chat_id
|
|
ORDER BY cc.total_cost_micros DESC;
|
|
|
|
-- name: GetChatCostPerUser :many
|
|
-- Deployment-wide per-user cost rollup within a date range.
|
|
-- Only counts assistant-role messages.
|
|
WITH chat_cost_users AS (
|
|
SELECT
|
|
c.owner_id AS user_id,
|
|
u.username,
|
|
u.name,
|
|
u.avatar_url,
|
|
COALESCE(SUM(cm.total_cost_micros), 0)::bigint AS total_cost_micros,
|
|
COUNT(*) FILTER (
|
|
WHERE cm.input_tokens IS NOT NULL
|
|
OR cm.output_tokens IS NOT NULL
|
|
OR cm.reasoning_tokens IS NOT NULL
|
|
OR cm.cache_creation_tokens IS NOT NULL
|
|
OR cm.cache_read_tokens IS NOT NULL
|
|
)::bigint AS message_count,
|
|
COUNT(DISTINCT COALESCE(c.root_chat_id, c.id))::bigint AS chat_count,
|
|
COALESCE(SUM(cm.input_tokens), 0)::bigint AS total_input_tokens,
|
|
COALESCE(SUM(cm.output_tokens), 0)::bigint AS total_output_tokens,
|
|
COALESCE(SUM(cm.cache_read_tokens), 0)::bigint AS total_cache_read_tokens,
|
|
COALESCE(SUM(cm.cache_creation_tokens), 0)::bigint AS total_cache_creation_tokens
|
|
FROM
|
|
chat_messages cm
|
|
JOIN
|
|
chats c ON c.id = cm.chat_id
|
|
JOIN
|
|
users u ON u.id = c.owner_id
|
|
WHERE
|
|
cm.role = 'assistant'
|
|
AND cm.created_at >= @start_date::timestamptz
|
|
AND cm.created_at < @end_date::timestamptz
|
|
AND (
|
|
@username::text = ''
|
|
OR u.username ILIKE '%' || @username::text || '%'
|
|
)
|
|
GROUP BY
|
|
c.owner_id,
|
|
u.username,
|
|
u.name,
|
|
u.avatar_url
|
|
)
|
|
SELECT
|
|
user_id,
|
|
username,
|
|
name,
|
|
avatar_url,
|
|
total_cost_micros,
|
|
message_count,
|
|
chat_count,
|
|
total_input_tokens,
|
|
total_output_tokens,
|
|
total_cache_read_tokens,
|
|
total_cache_creation_tokens,
|
|
COUNT(*) OVER()::bigint AS total_count
|
|
FROM
|
|
chat_cost_users
|
|
ORDER BY
|
|
total_cost_micros DESC,
|
|
username ASC
|
|
LIMIT
|
|
sqlc.arg('page_limit')::int
|
|
OFFSET
|
|
sqlc.arg('page_offset')::int;
|
|
|
|
-- name: GetChatUsageLimitConfig :one
|
|
SELECT * FROM chat_usage_limit_config WHERE singleton = TRUE LIMIT 1;
|
|
|
|
-- name: UpsertChatUsageLimitConfig :one
|
|
INSERT INTO chat_usage_limit_config (singleton, enabled, default_limit_micros, period, updated_at)
|
|
VALUES (TRUE, @enabled::boolean, @default_limit_micros::bigint, @period::text, NOW())
|
|
ON CONFLICT (singleton) DO UPDATE SET
|
|
enabled = EXCLUDED.enabled,
|
|
default_limit_micros = EXCLUDED.default_limit_micros,
|
|
period = EXCLUDED.period,
|
|
updated_at = NOW()
|
|
RETURNING *;
|
|
|
|
-- name: ListChatUsageLimitOverrides :many
|
|
SELECT u.id AS user_id, u.username, u.name, u.avatar_url,
|
|
u.chat_spend_limit_micros AS spend_limit_micros
|
|
FROM users u
|
|
WHERE u.chat_spend_limit_micros IS NOT NULL
|
|
ORDER BY u.username ASC;
|
|
|
|
-- name: UpsertChatUsageLimitUserOverride :one
|
|
UPDATE users
|
|
SET chat_spend_limit_micros = @spend_limit_micros::bigint
|
|
WHERE id = @user_id::uuid
|
|
RETURNING id AS user_id, username, name, avatar_url, chat_spend_limit_micros AS spend_limit_micros;
|
|
|
|
-- name: DeleteChatUsageLimitUserOverride :exec
|
|
UPDATE users SET chat_spend_limit_micros = NULL WHERE id = @user_id::uuid;
|
|
|
|
-- name: GetChatUsageLimitUserOverride :one
|
|
SELECT id AS user_id, chat_spend_limit_micros AS spend_limit_micros
|
|
FROM users
|
|
WHERE id = @user_id::uuid AND chat_spend_limit_micros IS NOT NULL;
|
|
|
|
-- name: GetUserChatSpendInPeriod :one
|
|
SELECT COALESCE(SUM(cm.total_cost_micros), 0)::bigint AS total_spend_micros
|
|
FROM chat_messages cm
|
|
JOIN chats c ON c.id = cm.chat_id
|
|
WHERE c.owner_id = @user_id::uuid
|
|
AND cm.created_at >= @start_time::timestamptz
|
|
AND cm.created_at < @end_time::timestamptz
|
|
AND cm.total_cost_micros IS NOT NULL;
|
|
|
|
-- name: CountEnabledModelsWithoutPricing :one
|
|
-- Counts enabled, non-deleted model configs that lack both input and
|
|
-- output pricing in their JSONB options.cost configuration.
|
|
SELECT COUNT(*)::bigint AS count
|
|
FROM chat_model_configs
|
|
WHERE enabled = TRUE
|
|
AND deleted = FALSE
|
|
AND (
|
|
options->'cost' IS NULL
|
|
OR options->'cost' = 'null'::jsonb
|
|
OR (
|
|
(options->'cost'->>'input_price_per_million_tokens' IS NULL)
|
|
AND (options->'cost'->>'output_price_per_million_tokens' IS NULL)
|
|
)
|
|
);
|
|
|
|
-- name: ListChatUsageLimitGroupOverrides :many
|
|
SELECT
|
|
g.id AS group_id,
|
|
g.name AS group_name,
|
|
g.display_name AS group_display_name,
|
|
g.avatar_url AS group_avatar_url,
|
|
g.chat_spend_limit_micros AS spend_limit_micros,
|
|
(SELECT COUNT(*)
|
|
FROM group_members_expanded gme
|
|
WHERE gme.group_id = g.id
|
|
AND gme.user_is_system = FALSE) AS member_count
|
|
FROM groups g
|
|
WHERE g.chat_spend_limit_micros IS NOT NULL
|
|
ORDER BY g.name ASC;
|
|
|
|
-- name: UpsertChatUsageLimitGroupOverride :one
|
|
UPDATE groups
|
|
SET chat_spend_limit_micros = @spend_limit_micros::bigint
|
|
WHERE id = @group_id::uuid
|
|
RETURNING id AS group_id, name, display_name, avatar_url, chat_spend_limit_micros AS spend_limit_micros;
|
|
|
|
-- name: DeleteChatUsageLimitGroupOverride :exec
|
|
UPDATE groups SET chat_spend_limit_micros = NULL WHERE id = @group_id::uuid;
|
|
|
|
-- name: GetChatUsageLimitGroupOverride :one
|
|
SELECT id AS group_id, chat_spend_limit_micros AS spend_limit_micros
|
|
FROM groups
|
|
WHERE id = @group_id::uuid AND chat_spend_limit_micros IS NOT NULL;
|
|
|
|
-- name: GetUserGroupSpendLimit :one
|
|
-- Returns the minimum (most restrictive) group limit for a user.
|
|
-- Returns -1 if the user has no group limits applied.
|
|
SELECT COALESCE(MIN(g.chat_spend_limit_micros), -1)::bigint AS limit_micros
|
|
FROM groups g
|
|
JOIN group_members_expanded gme ON gme.group_id = g.id
|
|
WHERE gme.user_id = @user_id::uuid
|
|
AND g.chat_spend_limit_micros IS NOT NULL;
|
|
|
|
-- name: ResolveUserChatSpendLimit :one
|
|
-- Resolves the effective spend limit for a user using the hierarchy:
|
|
-- 1. Individual user override (highest priority)
|
|
-- 2. Minimum group limit across all user's groups
|
|
-- 3. Global default from config
|
|
-- Returns -1 if limits are not enabled.
|
|
SELECT CASE
|
|
-- If limits are disabled, return -1.
|
|
WHEN NOT cfg.enabled THEN -1
|
|
-- Individual override takes priority.
|
|
WHEN u.chat_spend_limit_micros IS NOT NULL THEN u.chat_spend_limit_micros
|
|
-- Group limit (minimum across all user's groups) is next.
|
|
WHEN gl.limit_micros IS NOT NULL THEN gl.limit_micros
|
|
-- Fall back to global default.
|
|
ELSE cfg.default_limit_micros
|
|
END::bigint AS effective_limit_micros
|
|
FROM chat_usage_limit_config cfg
|
|
CROSS JOIN users u
|
|
LEFT JOIN LATERAL (
|
|
SELECT MIN(g.chat_spend_limit_micros) AS limit_micros
|
|
FROM groups g
|
|
JOIN group_members_expanded gme ON gme.group_id = g.id
|
|
WHERE gme.user_id = @user_id::uuid
|
|
AND g.chat_spend_limit_micros IS NOT NULL
|
|
) gl ON TRUE
|
|
WHERE u.id = @user_id::uuid
|
|
LIMIT 1;
|