jacob-mathison/query-orchestration
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Activity

Merged in feature/chatbot-limits (pull request #228)

Browse Source 
fix chatbot limits and test

* add tests
This commit is contained in:
Jay Brown
2026-05-28 19:23:39 +00:00
parent 451da3d26d
commit 33b0931b8c
6 changed files with 173 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/ai.generated/07-configuration-management.md
+2 -1
View File
@@ -137,7 +137,7 @@ From `internal/serviceconfig/chatbot/config.go`. Plan §8.
| `CHATBOT_SERVICE_URL` | (required, notEmpty) | env-loader rejects empty | FastAPI base URL. |
| `CHATBOT_API_KEY` | (required, notEmpty) | env-loader rejects empty | Sent verbatim as the `X-API-Key` header. |
| `CHATBOT_REQUEST_TIMEOUT_SECONDS` | `60` | `>= 1` | Per-attempt FastAPI HTTP timeout. |
| `CHATBOT_MAX_TURN_CHARS` | `2000` | `>= 1` | Caps prompt and `answer.text` rune count. |
| `CHATBOT_MAX_TURN_CHARS` | `2000` | `1 <= x <= 65536` | Caps prompt and `answer.text` rune count. Upper bound (`chatbot.MaxTurnCharsCeiling`) tracks the OpenAPI `prompt.maxLength` so the configured cap is always reachable. |
| `CHATBOT_RECENT_TURNS` | `5` | `1 <= x <= 5` | Conversation-history length and session-list preview cap. |
Validated at startup via `chatbot.ChatbotConfig.Validate()`, called
@@ -147,6 +147,7 @@ validation failure.
Derived constants (not env-driven):
- `chatbot.StateMaxBytes = 64 * 1024` — session-state cap on persist.
- `chatbot.MaxTurnCharsCeiling = 65536` — upper bound on `CHATBOT_MAX_TURN_CHARS`, mirroring the OpenAPI `BotTurnAddRequest.prompt.maxLength` in `serviceAPIs/queryAPI.yaml`. Changing one requires changing the other.
- `chatbot.ChatbotConfig.InflightGraceSeconds() = 2 * RequestTimeoutSeconds`
per plan §8 literal.
docs/ai.generated/12-chatbot-guide.md
+2 -1
View File
@@ -163,12 +163,13 @@ Five env vars, plan §8. Validated at startup via
| `CHATBOT_SERVICE_URL` | (required, notEmpty) | env-loader rejects empty | FastAPI base URL. |
| `CHATBOT_API_KEY` | (required, notEmpty) | env-loader rejects empty | Sent verbatim as `X-API-Key` header. |
| `CHATBOT_REQUEST_TIMEOUT_SECONDS` | 60 | `>= 1` | Per-attempt FastAPI timeout. |
| `CHATBOT_MAX_TURN_CHARS` | 2000 | `>= 1` | Caps prompt and `answer.text` rune count. |
| `CHATBOT_MAX_TURN_CHARS` | 2000 | `1 <= x <= 65536` | Caps prompt and `answer.text` rune count. Upper bound (`chatbot.MaxTurnCharsCeiling`) mirrors the OpenAPI `prompt.maxLength`. |
| `CHATBOT_RECENT_TURNS` | 5 | `1 <= x <= 5` | Conversation-history length and session-list preview cap. |
Derived values (constants/helpers, not env vars):
- `chatbot.StateMaxBytes = 64 * 1024` — session-state cap (plan §8).
- `chatbot.MaxTurnCharsCeiling = 65536` — upper bound for `CHATBOT_MAX_TURN_CHARS`, locked to the OpenAPI `BotTurnAddRequest.prompt.maxLength`. Validate() rejects any value above this.
- `chatbot.ChatbotConfig.InflightGraceSeconds() = 2 * RequestTimeoutSeconds`
per plan §8 literal. The service uses
`effectiveGraceSeconds(cfg) = max(2 * RequestTimeoutSeconds, 60)` as
internal/serviceconfig/chatbot/config.go
+15
View File
@@ -15,6 +15,16 @@ import "fmt"
// 64 KiB matches the plan literal `64 * 1024`.
const StateMaxBytes = 64 * 1024
// MaxTurnCharsCeiling is the largest legal value for MaxTurnChars.
// It mirrors the OpenAPI schema constraint on the prompt body at
// serviceAPIs/queryAPI.yaml (BotTurnAddRequest.prompt.maxLength=65536).
// Setting CHATBOT_MAX_TURN_CHARS above this ceiling would create a
// silent inconsistency: the schema validator would reject oversized
// prompts before the service-layer length guard ran, leaving operators
// with no signal that the configured cap is unreachable. Any change to
// this value must be made in lockstep with the OpenAPI spec.
const MaxTurnCharsCeiling = 65536
// ChatbotConfig is the queryAPI-side configuration block for the FastAPI
// chatbot agent service. Fields, env-var names, defaults, and the
// notEmpty/required tags are dictated verbatim by plan §8.
@@ -53,6 +63,11 @@ func (c ChatbotConfig) Validate() error {
if c.MaxTurnChars < 1 {
return fmt.Errorf("chatbot config: MaxTurnChars must be >= 1, got %d", c.MaxTurnChars)
}
if c.MaxTurnChars > MaxTurnCharsCeiling {
return fmt.Errorf(
"chatbot config: MaxTurnChars must be <= %d (OpenAPI prompt.maxLength), got %d",
MaxTurnCharsCeiling, c.MaxTurnChars)
}
if c.RecentTurns < 1 || c.RecentTurns > 5 {
return fmt.Errorf("chatbot config: RecentTurns must be in [1, 5], got %d", c.RecentTurns)
}
internal/serviceconfig/chatbot/config_test.go
+90
View File
@@ -266,6 +266,96 @@ func TestChatbotConfig_Validate(t *testing.T) {
}
}
// TestChatbotConfig_MaxTurnChars_NonDefault_Loaded covers Gap 1: prior
// to this test no case asserted that CHATBOT_MAX_TURN_CHARS=5000 (a
// non-default value) was both loaded by env.ParseWithOptions and
// accepted by Validate(). The plan permits any operator-chosen value
// >= 1; this test pins that contract for a representative larger cap
// so a regression that hard-codes the default cannot pass silently.
func TestChatbotConfig_MaxTurnChars_NonDefault_Loaded(t *testing.T) {
t.Parallel()
envMap := requiredEnvBase()
envMap["CHATBOT_MAX_TURN_CHARS"] = "5000"
cfg, err := loadConfig(t, envMap)
require.NoError(t, err, "loader must accept CHATBOT_MAX_TURN_CHARS=5000")
require.Equal(t, 5000, cfg.MaxTurnChars,
"non-default CHATBOT_MAX_TURN_CHARS must be honored by the loader")
require.NoError(t, cfg.Validate(),
"Validate must accept a non-default MaxTurnChars within the legal range")
}
// TestChatbotConfig_MaxTurnChars_AtOpenAPILimit_Accepted locks the
// upper-edge boundary: MaxTurnChars=65536 (exactly the OpenAPI
// prompt.maxLength) must load and validate. Pinned alongside the
// rejection cases so an off-by-one in the ceiling check fails loudly.
func TestChatbotConfig_MaxTurnChars_AtOpenAPILimit_Accepted(t *testing.T) {
t.Parallel()
envMap := requiredEnvBase()
envMap["CHATBOT_MAX_TURN_CHARS"] = "65536"
cfg, err := loadConfig(t, envMap)
require.NoError(t, err, "loader must accept CHATBOT_MAX_TURN_CHARS=65536")
require.Equal(t, chatbot.MaxTurnCharsCeiling, cfg.MaxTurnChars,
"loaded value must equal the documented ceiling")
require.NoError(t, cfg.Validate(),
"Validate must accept MaxTurnChars exactly at MaxTurnCharsCeiling")
}
// TestChatbotConfig_MaxTurnChars_AboveOpenAPILimit_Rejected covers Gap 2:
// the OpenAPI schema (serviceAPIs/queryAPI.yaml: prompt.maxLength=65536)
// caps the prompt body at 65,536 runes. If CHATBOT_MAX_TURN_CHARS is set
// above that ceiling the system enters an inconsistent state — the
// generated request validator silently truncates the cap to 65,536 while
// the service-layer guard believes a higher cap is in force. Validate()
// MUST refuse this configuration so the inconsistency is caught at boot,
// not on the first oversized prompt. Expected to FAIL until backend-eng
// adds an upper-bound check tied to the OpenAPI maxLength.
func TestChatbotConfig_MaxTurnChars_AboveOpenAPILimit_Rejected(t *testing.T) {
t.Parallel()
envMap := requiredEnvBase()
envMap["CHATBOT_MAX_TURN_CHARS"] = "65537"
cfg, err := loadConfig(t, envMap)
require.NoError(t, err, "loader must parse the integer value")
validateErr := cfg.Validate()
require.Error(t, validateErr,
"Validate must reject MaxTurnChars > OpenAPI prompt.maxLength=65536; "+
"otherwise oversized prompts are silently rejected by the schema validator "+
"with no operator-visible signal that the configured cap is unreachable")
require.True(t,
strings.Contains(validateErr.Error(), "MaxTurnChars"),
"error %q must name the offending field for actionable diagnostics",
validateErr.Error())
}
// TestChatbotConfig_MaxTurnChars_GrosslyOverLimit_Rejected covers Gap 3:
// Validate() currently only enforces MaxTurnChars >= 1 — there is no
// upper bound at all. An operator typo such as CHATBOT_MAX_TURN_CHARS=
// 10000000 loads cleanly today, silently exceeds the OpenAPI ceiling,
// and creates the inconsistency described in Gap 2 at scale. This test
// asserts that any value above the documented ceiling is rejected.
// Expected to FAIL until backend-eng adds the upper-bound check.
func TestChatbotConfig_MaxTurnChars_GrosslyOverLimit_Rejected(t *testing.T) {
t.Parallel()
envMap := requiredEnvBase()
envMap["CHATBOT_MAX_TURN_CHARS"] = "10000000"
cfg, err := loadConfig(t, envMap)
require.NoError(t, err, "loader must parse the integer value")
validateErr := cfg.Validate()
require.Error(t, validateErr,
"Validate must reject grossly over-limit MaxTurnChars to fail-fast on operator typos")
require.True(t,
strings.Contains(validateErr.Error(), "MaxTurnChars"),
"error %q must name the offending field",
validateErr.Error())
}
// TestChatbotConfig_InflightGraceSeconds locks down the derived value
// from plan §8: "InflightGraceSeconds = 2 * RequestTimeoutSeconds". The
// helper is consumed by AddTurn (M4) to size the grace window passed to
scripts/manual_tests/test_two_pool_auth.py
+35 -20
View File
@@ -2,7 +2,7 @@
"""End-to-end manual test for the two-pool service-account auth path.
Validates plans/api.key.support.cognito.two-pool.plan.3.claude.md
against a localhost queryAPI configured with both pool A
against a running queryAPI configured with both pool A
(COGNITO_USER_POOL_ID) and pool B (COGNITO_SVC_USER_POOL_ID), and
auth enabled.
@@ -14,38 +14,42 @@ What this script does (top to bottom):
Step 3. Decode the AccessToken locally (no signature check) to
confirm sub, iss, token_use match the values from the
provisioning log.
Step 4. Call GET /admin/users with Authorization: Bearer
<access_token>. /admin/users is a super_admin-only
Step 4. Call GET /admin/eula with Authorization: Bearer
<access_token>. /admin/eula is a super_admin-only
endpoint, so a 200 proves: (a) queryAPI's multi-issuer
verifier accepted the pool B token; (b) the service
account was successfully registered in Permit with the
super_admin role; (c) the env-check dispatched to pool B.
Step 5. Call GET /admin/eula with the same token. Second
super_admin-only endpoint — covers the "pick at least 2"
requirement from the goal statement.
Step 6. Print a summary.
account has super_admin in Permit; (c) the env-check
dispatched to pool B.
Step 5. Call GET /admin/eula/agreements with the same token.
Second super_admin-only endpoint — covers the "pick at
least 2" requirement from the goal statement.
Step 6. Print PASS/FAIL summary.
Zero third-party deps — stdlib only, like
cmd/auth_related/service.accounts/test.scripts/test.service.account.auth.py.
Prerequisites:
- queryAPI is running locally on http://localhost:8080.
- queryAPI is reachable at --base-url (default http://localhost:8080).
- queryAPI was started with both COGNITO_USER_POOL_ID and
COGNITO_SVC_USER_POOL_ID configured.
- DISABLE_AUTH is NOT set (we want real Cognito JWT validation).
- The test service account has been registered in Permit.io and
assigned super_admin via:
./cmd/auth_related/service.account.tool/run.tool.sh \
register-permit svc-test-admin-1 --roles super_admin
- SVC_TEST_PASSWORD is supplied via env (see .env.example). Any
.env between this file and the filesystem root is auto-loaded.
Everything else (pool ID, app client ID, username, sub, region)
is hardcoded below — public identifiers, not secrets — so the
script "just works" without further shell setup.
- The test service account already has super_admin in the Permit
instance backing the target queryAPI. This script does not
provision Permit — that is an environment-setup concern.
- SVC_TEST_PASSWORD is supplied via env (the repo-root .env.example
documents the variable). Any .env between this file and the
filesystem root is auto-loaded. Everything else (pool ID, app
client ID, username, sub, region) is hardcoded below — public
identifiers, not secrets — so the script "just works" without
further shell setup.
- INSECURE_TLS=1 disables queryAPI TLS cert verification (Cognito
is left alone). Needed when queryAPI is fronted by an ALB with
a self-signed cert. Off by default.
Usage:
python3 scripts/manual_tests/test_two_pool_auth.py
python3 scripts/manual_tests/test_two_pool_auth.py --base-url http://localhost:8080
INSECURE_TLS=1 python3 scripts/manual_tests/test_two_pool_auth.py --base-url https://...
"""
from __future__ import annotations
@@ -54,6 +58,7 @@ import argparse
import base64
import json
import os
import ssl
import sys
import urllib.error
import urllib.request
@@ -122,6 +127,13 @@ for key in ("AWS_ENDPOINT_URL", "AWS_ENDPOINT_URL_COGNITO_IDP"):
PASSWORD = _required("SVC_TEST_PASSWORD")
# INSECURE_TLS=1 disables certificate verification for queryAPI calls
# only (Cognito is left alone — its cert is real). Needed when queryAPI
# is fronted by an ALB with a self-signed cert. Off by default so the
# bypass cannot slip in silently.
INSECURE_TLS = os.environ.get("INSECURE_TLS", "").strip().lower() in ("1", "true", "yes", "y")
_TLS_CTX = ssl._create_unverified_context() if INSECURE_TLS else None
ISSUER_URL = f"https://cognito-idp.{AWS_REGION}.amazonaws.com/{USER_POOL_ID}"
COGNITO_IDP_URL = f"https://cognito-idp.{AWS_REGION}.amazonaws.com/"
@@ -180,7 +192,7 @@ def queryapi_get(base_url: str, path: str, access_token: str = "") -> tuple[int,
headers["Authorization"] = "Bearer " + access_token
req = urllib.request.Request(url, headers=headers, method="GET")
try:
with urllib.request.urlopen(req, timeout=HTTP_TIMEOUT) as r:
with urllib.request.urlopen(req, timeout=HTTP_TIMEOUT, context=_TLS_CTX) as r:
return r.status, r.read().decode("utf-8", errors="replace")
except urllib.error.HTTPError as e:
body = e.read().decode("utf-8", errors="replace")
@@ -193,6 +205,9 @@ def main() -> int:
args = parser.parse_args()
base = args.base_url
if INSECURE_TLS:
print("WARNING: INSECURE_TLS=1 — queryAPI TLS verification is DISABLED.")
section("Step 1: queryAPI health")
code, body = queryapi_get(base, "/health")
print(f"GET {base}/health -> {code}")
scripts/manual_tests/test_two_pool_auth.sh
+28 -32
View File
@@ -1,51 +1,47 @@
#!/usr/bin/env bash
# test_two_pool_auth.sh runs the end-to-end integration test for the
# two-pool service-account auth path. It assumes the operator has:
# - AWS credentials with read access to the test pools
# (us-east-2_<pool-a>, us-east-2_Qom1i9qIG)
# - PERMIT_IO_API_KEY set in the environment
# - Docker running (for `task compose:refresh`)
# two-pool service-account auth path against a running queryAPI.
#
# The script:
# 1. Registers the pre-existing test service-account in Permit with
# super_admin (idempotent — re-running is safe).
# 2. Reminds the operator to boot queryAPI with both pool A and B
# configured (compose:refresh by itself does not know about the
# new env vars; the operator must source the right .env).
# 3. Runs scripts/manual_tests/test_two_pool_auth.py against the
# local queryAPI to validate the end-to-end flow.
# The script assumes the target environment is already provisioned:
# - The test service account exists in the Cognito service pool
# (us-east-2_Qom1i9qIG).
# - It is already registered in the Permit instance backing the
# target queryAPI with the super_admin role. This script does NOT
# touch Permit — that is an environment-setup concern.
# - SVC_TEST_PASSWORD is available (in env or any .env walked up
# from this file to the filesystem root).
#
# Usage:
# PERMIT_IO_API_KEY=permit_key_... \
# AWS_PROFILE=aarete \
# ./scripts/manual_tests/test_two_pool_auth.sh
# BASE_URL=http://localhost:8080 ./scripts/manual_tests/test_two_pool_auth.sh
# INSECURE_TLS=0 ./scripts/manual_tests/test_two_pool_auth.sh # require valid cert
set -euo pipefail
if [[ -z "${PERMIT_IO_API_KEY:-}" ]]; then
echo "ERROR: PERMIT_IO_API_KEY must be set" >&2
exit 1
fi
# ---------------------------------------------------------------------------
# Target queryAPI to test against. Uncomment ONE line below.
# Override at the CLI with: BASE_URL=https://... ./test_two_pool_auth.sh
# ---------------------------------------------------------------------------
: "${BASE_URL:=https://queryo-query-q0pz6j2syaox-1001614225.us-east-2.elb.amazonaws.com}"
# : "${BASE_URL:=http://localhost:8080}"
# ---------------------------------------------------------------------------
# The default BASE_URL above is an ALB with a self-signed cert. Python's
# urllib rejects that by default, so opt into the bypass here. Unset or
# set to 0 if you point BASE_URL at a host with a real CA-issued cert.
: "${INSECURE_TLS:=1}"
export INSECURE_TLS
REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
echo "Step 1/3: registering svc-test-admin-1 in Permit with super_admin..."
pushd "$REPO_ROOT/cmd/auth_related/service.account.tool" >/dev/null
COGNITO_REGION="us-east-2" \
COGNITO_SVC_USER_POOL_ID="us-east-2_Qom1i9qIG" \
COGNITO_SVC_APP_CLIENT_ID="6r7chkvjotgs2bsfbd5ibaaft7" \
go run ./... register-permit svc-test-admin-1 --roles super_admin
popd >/dev/null
echo
echo "Step 2/3: confirm queryAPI is running locally with pool B enabled."
echo "Step 1/2: confirm queryAPI is running with pool B enabled."
echo " - COGNITO_USER_POOL_ID should be set (existing pool A)"
echo " - COGNITO_SVC_USER_POOL_ID should be us-east-2_Qom1i9qIG (pool B)"
echo " - DISABLE_AUTH must be unset (we are testing real auth)"
echo " - queryAPI must be reachable at http://localhost:8080"
echo " - queryAPI must be reachable at ${BASE_URL}"
echo
read -rp "Press enter when queryAPI is up and configured..."
echo
echo "Step 3/3: running test_two_pool_auth.py against localhost..."
python3 "$REPO_ROOT/scripts/manual_tests/test_two_pool_auth.py"
echo "Step 2/2: running test_two_pool_auth.py against ${BASE_URL}..."
python3 "$REPO_ROOT/scripts/manual_tests/test_two_pool_auth.py" --base-url "${BASE_URL}"