Merge branch 'main' into main
This commit is contained in:
commit
803144adaf
|
|
@ -9,6 +9,9 @@
|
|||
# =============================================================================
|
||||
LOG_LEVEL=INFO
|
||||
PERFORMANCE_LOG_FORMAT=compact # compact|rich
|
||||
# API server processes used by the Docker entrypoint (default: 1).
|
||||
# Each process owns a separate pool when connection pooling is enabled.
|
||||
# API_WORKERS=1
|
||||
# SESSION_OBSERVERS_LIMIT=10
|
||||
# GET_CONTEXT_MAX_TOKENS=100000
|
||||
# MAX_FILE_SIZE=5242880 # Bytes
|
||||
|
|
|
|||
15
CHANGELOG.md
15
CHANGELOG.md
|
|
@ -11,6 +11,21 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
|
|||
|
||||
- Qdrant vector store backend (`VECTOR_STORE_TYPE=qdrant`) as an optional `qdrant` extra (#683)
|
||||
|
||||
## [3.1.1] - 2026-09-02
|
||||
|
||||
### Changed
|
||||
|
||||
- Server `requires-python` is `>=3.13`, matching the production image. Self-hosters on 3.10–3.12 need to upgrade; SDK and CLI floors are unchanged (#1090)
|
||||
|
||||
### Fixed
|
||||
|
||||
- Concurrent `create_documents` writers to the same collection deadlocked on `times_derived` reinforcement UPDATEs issued in batch order; the error was swallowed per-document, the batch was lost, and the queue item was marked processed. Writers now lock target rows with `SELECT ... ORDER BY id FOR UPDATE` before applying, abort the batch on `SQLAlchemyError` instead of continuing through a dead session, and retry transient errors (deadlock, serialization failure, lock/statement timeout, lost connection) up to `MAX_RETRYABLE_ATTEMPTS` instead of burning the item (#1033)
|
||||
- Scope backfill no longer embeds, writes, and syncs every planned copy at once. A 14k-document session is ~580MB of vectors; several concurrent backfills OOM-killed the deriver at its 1000Mi limit and crash-looped because the work units never completed. Phases 2–4 now run per chunk of 500 specs, reload source embeddings per chunk, and drop them once synced. Membership is locked across chunk writes so a concurrent leave cannot commit between the check and the inserts (#1104)
|
||||
- Model-generated observations with NUL bytes (`\u0000`) no longer fail the exact-content dedup pre-fetch with a Postgres `DataError` that dropped the whole observer batch. Ingress already stripped NUL from user content; the deriver now strips it so stored text matches embedded text. All-NUL content is dropped rather than stored empty (#1095)
|
||||
- `search_messages` no longer forwards `top_k=0` to Turbopuffer (which requires 1..10000). Zero/negative limits short-circuit to empty results; tool limits are floored at 1. The documents path was already guarded (#970); this closes the message path (#1084)
|
||||
- OpenAI-compatible tool-call turns with `content=null` keep null through history replay instead of being coerced to `""`. Providers that bind reasoning state to the exact assistant message shape were breaking on the empty string. Tool-less null still becomes `""` (#1064)
|
||||
- The production image now ships `pyproject.toml` in the runtime stage, so the service reports its real version instead of `unknown` in OpenAPI and telemetry (#1074)
|
||||
|
||||
## [3.1.0] - 2026-08-25
|
||||
|
||||
### Added
|
||||
|
|
|
|||
10
CLAUDE.md
10
CLAUDE.md
|
|
@ -118,6 +118,16 @@ cd sdks/typescript && bun run tsc --noEmit
|
|||
- **Never hold a DB session during external calls** (LLM, embedding, HTTP). If a function needs both a DB session and an external call result, compute the external result first and pass it as a parameter. This avoids tying up DB connections during slow network I/O. Use `tracked_db` for short-lived, DB-only operations; pass a shared session when multiple DB-only calls can reuse one connection.
|
||||
- **Never write through a read-only session** (`tracked_db(..., read_only=True)`, `get_read_db`, `ReadSessionLocal`). These run in AUTOCOMMIT mode with no transaction: writes are NOT blocked by the database — they silently commit immediately, and `begin_nested()` savepoints break. There is no runtime guard; this is enforced by convention only. Use `read_only=True` strictly for SELECT-only windows; anything that mutates (including get-or-create paths) must use a regular write session.
|
||||
|
||||
#### Multi-row locking and deadlocks
|
||||
|
||||
Tables written concurrently by more than one worker — `documents` (deriver, dreamer, scope backfill/removal, reconciler) and `queue` (every deriver replica) — deadlock when two writers touch an overlapping row set in different orders. Rules:
|
||||
|
||||
- **A multi-row `SELECT ... FOR UPDATE` MUST carry an explicit `ORDER BY <pk>`.** Without it Postgres locks in scan order, which differs per plan, so two writers with overlapping sets can cycle. `_apply_document_row_updates` in `src/crud/document.py` is the reference implementation.
|
||||
- **`WHERE id IN (...)` does NOT impose an order**, so sorting the Python list is a no-op — the list order is discarded and the planner picks `Bitmap Heap Scan` (ctid order), `Index Scan` (id order), or `Seq Scan` per invocation. Deterministic ordering requires either a preceding `SELECT ... ORDER BY id FOR UPDATE` or `WHERE id IN (SELECT id ... ORDER BY id FOR UPDATE)`.
|
||||
- **`Document.id` is a random nanoid** (`models.py`), so id order is uncorrelated with physical order — an unordered predicate `UPDATE`/`DELETE` is roughly a coin flip against an id-ordered locker per row pair, not a rare edge case. (`QueueItem.id` is an integer identity, so there id order is also chronological.)
|
||||
- **Prefer no lock at all.** A single `UPDATE ... WHERE <predicate>` acquires row locks as it writes and has no separate lock phase to get wrong. Reach for `FOR UPDATE` only when a value must be read, computed in Python, and written back — that read-modify-write is the only reason `_apply_document_row_updates` locks (it replaced a server-side `func.greatest()`), and `populate_existing=True` is required with it so the identity map doesn't serve a stale pre-lock value. Server-side expressions (`func.greatest`, the JSONB `-` operator) avoid the lock entirely; see `_clear_work_unit_retry_attempts` in `src/deriver/queue_manager.py`.
|
||||
- `FOR UPDATE SKIP LOCKED` (the reconciler's claim pattern) never waits, so it cannot be a deadlock partner — but holding those locks across an external call still stalls other writers. See the "never hold a DB session during external calls" rule above.
|
||||
|
||||
#### Auth scoping
|
||||
|
||||
- **`allow_member_read=True` (in `require_auth(...)`) is read-only — NEVER set it on a route that mutates state.** It lets a peer-scoped key reach a session route when its peer is an active member of the session, so on a mutating route it would hand any session member write access (message injection, config mutation, deletion). HTTP method is not a reliable read/write signal here (some read routes use POST for a richer body), so this is enforced by an explicit allowlist in `tests/routes/test_auth_route_policy.py` — adding the flag to a new route fails that test until you consciously add the route to `EXPECTED_MEMBER_READ_ROUTES`, and you must never add a mutating method there.
|
||||
|
|
|
|||
51
README.md
51
README.md
|
|
@ -8,7 +8,7 @@
|
|||
|
||||
---
|
||||
|
||||

|
||||

|
||||
[](https://pypi.org/project/honcho-ai/)
|
||||
[](https://npmjs.org/package/@honcho-ai/sdk)
|
||||
[](https://pypi.org/project/honcho-cli/)
|
||||
|
|
@ -173,6 +173,23 @@ See the full [SDK Reference](https://honcho.dev/docs/v3/documentation/reference/
|
|||
|
||||
## Integrations
|
||||
|
||||
Honcho ships a first-party memory plugin for every major coding agent. They all read the same
|
||||
`~/.honcho/config.json`, so one key configures all of them — and pointing two at the same `workspace`
|
||||
gives them one shared memory.
|
||||
|
||||
| Agent | Install | Source |
|
||||
| ---------------- | ------------------------------------------------------- | ------------------------------------------------------------------ |
|
||||
| Claude Code | `/plugin marketplace add plastic-labs/claude-honcho` | [claude-honcho](https://github.com/plastic-labs/claude-honcho) |
|
||||
| Codex | `npm install -g @honcho-ai/codex-honcho` | [codex-honcho](https://github.com/plastic-labs/codex-honcho) |
|
||||
| Cursor | `curl -fsSL .../cursor-honcho/main/install.sh \| bash` | [cursor-honcho](https://github.com/plastic-labs/cursor-honcho) |
|
||||
| DeepSeek Harness | `dsh plugin --profile <name> add @honcho-ai/dsh-honcho` | [dsh-honcho](https://github.com/plastic-labs/dsh-honcho) |
|
||||
| OpenCode | `opencode plugin "@honcho-ai/opencode-honcho" --global` | [opencode-honcho](https://github.com/plastic-labs/opencode-honcho) |
|
||||
| OpenClaw | `openclaw plugins install @honcho-ai/openclaw-honcho` | [openclaw-honcho](https://github.com/plastic-labs/openclaw-honcho) |
|
||||
| Hermes | `hermes memory setup` | built in upstream |
|
||||
| Any MCP client | `claude mcp add honcho --transport http ...` | [MCP guide](https://honcho.dev/docs/v3/guides/integrations/mcp) |
|
||||
|
||||
Get a key at [app.honcho.dev](https://app.honcho.dev), then `honcho init` (or `uv tool install honcho-cli && honcho init`) writes it to `~/.honcho/config.json` once for every integration.
|
||||
|
||||
### Claude Code
|
||||
|
||||
Two ways, depending on how deep you want to go:
|
||||
|
|
@ -194,7 +211,33 @@ claude mcp add honcho \
|
|||
--header "X-Honcho-User-Name: YourName"
|
||||
```
|
||||
|
||||
Details: [Claude Code guide](https://honcho.dev/docs/v3/guides/integrations/claude-code) · [MCP guide](https://honcho.dev/docs/v3/guides/integrations/mcp).
|
||||
Details: [Claude Code guide](https://honcho.dev/docs/v3/guides/integrations/claude-code) · [MCP guide](https://honcho.dev/docs/v3/guides/integrations/mcp) · [repo](https://github.com/plastic-labs/claude-honcho).
|
||||
|
||||
### Codex
|
||||
|
||||
```bash
|
||||
npm install -g @honcho-ai/codex-honcho
|
||||
codex-honcho install # registers hooks + MCP + skill in ~/.codex
|
||||
```
|
||||
|
||||
Restart Codex to load the hooks. Details: [Codex guide](https://honcho.dev/docs/v3/guides/integrations/codex) · [repo](https://github.com/plastic-labs/codex-honcho).
|
||||
|
||||
### Cursor
|
||||
|
||||
```bash
|
||||
curl -fsSL https://raw.githubusercontent.com/plastic-labs/cursor-honcho/main/install.sh | bash
|
||||
```
|
||||
|
||||
Windows (PowerShell): `irm https://raw.githubusercontent.com/plastic-labs/cursor-honcho/main/install.ps1 | iex`. The installer wires global hooks and MCP config. Details: [cursor-honcho](https://github.com/plastic-labs/cursor-honcho).
|
||||
|
||||
### DeepSeek Harness
|
||||
|
||||
```bash
|
||||
dsh plugin --profile <name> add @honcho-ai/dsh-honcho
|
||||
```
|
||||
|
||||
A native Cordis plugin. It injects memory into the system prompt and captures new information from the session event feed. The model gets three tools — honcho_search, honcho_chat, and honcho_remember — and you can run /honcho to check status.
|
||||
Details: [DeepSeek Harness guide](https://honcho.dev/docs/v3/guides/integrations/deepseek-harness) · [repo](https://github.com/plastic-labs/dsh-honcho).
|
||||
|
||||
### OpenCode
|
||||
|
||||
|
|
@ -202,7 +245,7 @@ Details: [Claude Code guide](https://honcho.dev/docs/v3/guides/integrations/clau
|
|||
opencode plugin "@honcho-ai/opencode-honcho" --global
|
||||
```
|
||||
|
||||
Details: [OpenCode guide](https://honcho.dev/docs/v3/guides/integrations/opencode).
|
||||
Details: [OpenCode guide](https://honcho.dev/docs/v3/guides/integrations/opencode) · [repo](https://github.com/plastic-labs/opencode-honcho).
|
||||
|
||||
### OpenClaw
|
||||
|
||||
|
|
@ -212,7 +255,7 @@ openclaw honcho setup
|
|||
openclaw gateway --force
|
||||
```
|
||||
|
||||
`openclaw honcho setup` prompts for your API key, writes the config, and optionally migrates legacy `MEMORY.md` / `USER.md` / `IDENTITY.md` files into Honcho (non-destructive — originals are never deleted). Details: [OpenClaw guide](https://honcho.dev/docs/v3/guides/integrations/openclaw).
|
||||
`openclaw honcho setup` prompts for your API key, writes the config, and optionally migrates legacy `MEMORY.md` / `USER.md` / `IDENTITY.md` files into Honcho (non-destructive — originals are never deleted). Details: [OpenClaw guide](https://honcho.dev/docs/v3/guides/integrations/openclaw) · [repo](https://github.com/plastic-labs/openclaw-honcho).
|
||||
|
||||
### Hermes
|
||||
|
||||
|
|
|
|||
|
|
@ -80,6 +80,34 @@ services:
|
|||
required: false
|
||||
restart: unless-stopped
|
||||
|
||||
mcp:
|
||||
build:
|
||||
context: ./mcp
|
||||
dockerfile: Dockerfile
|
||||
depends_on:
|
||||
api:
|
||||
condition: service_healthy
|
||||
ports:
|
||||
- "127.0.0.1:3000:3000"
|
||||
environment:
|
||||
- HONCHO_API_URL=http://api:8000
|
||||
env_file:
|
||||
- path: .env
|
||||
required: false
|
||||
healthcheck:
|
||||
test:
|
||||
[
|
||||
"CMD",
|
||||
"bun",
|
||||
"-e",
|
||||
"fetch('http://127.0.0.1:3000/health').then((r)=>process.exit(r.ok?0:1)).catch(()=>process.exit(1))",
|
||||
]
|
||||
interval: 5s
|
||||
timeout: 5s
|
||||
retries: 5
|
||||
start_period: 10s
|
||||
restart: unless-stopped
|
||||
|
||||
database:
|
||||
image: pgvector/pgvector:pg15
|
||||
restart: unless-stopped
|
||||
|
|
|
|||
|
|
@ -5,4 +5,4 @@ echo "Running database migrations..."
|
|||
/app/.venv/bin/python scripts/provision_db.py
|
||||
|
||||
echo "Starting API server..."
|
||||
exec /app/.venv/bin/fastapi run --host 0.0.0.0 src/main.py
|
||||
exec /app/.venv/bin/fastapi run --host 0.0.0.0 --workers "${API_WORKERS:-1}" src/main.py
|
||||
|
|
|
|||
|
|
@ -30,7 +30,8 @@ This guide helps you match the right SDK version to your Honcho API version. New
|
|||
|
||||
| Honcho API Version | TypeScript SDK | Python SDK |
|
||||
|-------------------|---------------|------------|
|
||||
| v3.1.0 (Current) | v2.4.0 | v2.4.0 |
|
||||
| v3.1.1 (Current) | v2.4.0 | v2.4.0 |
|
||||
| v3.1.0 | v2.4.0 | v2.4.0 |
|
||||
| v3.0.12 | v2.3.0 | v2.3.0 |
|
||||
| v3.0.11 | v2.1.2 | v2.1.2 |
|
||||
| v3.0.10 | v2.1.2 | v2.1.2 |
|
||||
|
|
|
|||
|
|
@ -27,7 +27,22 @@ Welcome to the Honcho changelog! This section documents all notable changes to t
|
|||
### Honcho API and SDK Changelogs
|
||||
<Tabs>
|
||||
<Tab title="Honcho API">
|
||||
<Update label="v3.1.0 (Current)">
|
||||
<Update label="v3.1.1 (Current)">
|
||||
### Changed
|
||||
|
||||
- Server `requires-python` is `>=3.13`, matching the production image. Self-hosters on 3.10–3.12 need to upgrade; SDK and CLI floors are unchanged (#1090)
|
||||
|
||||
### Fixed
|
||||
|
||||
- Concurrent `create_documents` writers to the same collection deadlocked on `times_derived` reinforcement UPDATEs issued in batch order; the error was swallowed per-document, the batch was lost, and the queue item was marked processed. Writers now lock target rows with `SELECT ... ORDER BY id FOR UPDATE` before applying, abort the batch on `SQLAlchemyError` instead of continuing through a dead session, and retry transient errors (deadlock, serialization failure, lock/statement timeout, lost connection) up to `MAX_RETRYABLE_ATTEMPTS` instead of burning the item (#1033)
|
||||
- Scope backfill no longer embeds, writes, and syncs every planned copy at once. A 14k-document session is ~580MB of vectors; several concurrent backfills OOM-killed the deriver at its 1000Mi limit and crash-looped because the work units never completed. Phases 2–4 now run per chunk of 500 specs, reload source embeddings per chunk, and drop them once synced. Membership is locked across chunk writes so a concurrent leave cannot commit between the check and the inserts (#1104)
|
||||
- Model-generated observations with NUL bytes (`\u0000`) no longer fail the exact-content dedup pre-fetch with a Postgres `DataError` that dropped the whole observer batch. Ingress already stripped NUL from user content; the deriver now strips it so stored text matches embedded text. All-NUL content is dropped rather than stored empty (#1095)
|
||||
- `search_messages` no longer forwards `top_k=0` to Turbopuffer (which requires 1..10000). Zero/negative limits short-circuit to empty results; tool limits are floored at 1. The documents path was already guarded (#970); this closes the message path (#1084)
|
||||
- OpenAI-compatible tool-call turns with `content=null` keep null through history replay instead of being coerced to `""`. Providers that bind reasoning state to the exact assistant message shape were breaking on the empty string. Tool-less null still becomes `""` (#1064)
|
||||
- The production image now ships `pyproject.toml` in the runtime stage, so the service reports its real version instead of `unknown` in OpenAPI and telemetry (#1074)
|
||||
</Update>
|
||||
|
||||
<Update label="v3.1.0">
|
||||
### Added
|
||||
|
||||
- Scopes: a named grouping of sessions that acts as a visibility boundary on recall, implemented as a facade over an observer peer (`scope.{name}` with `{"kind": "scope"}`). Developers manage them exclusively through `/v3/workspaces/{workspace_id}/scopes` (create-or-get, list, get, add/list/remove session membership) and an optional `scopes` field on session create — never through the observer/observed mechanics. Scope peers cannot author messages, cannot be a chat or representation `target`, are excluded from `peers.list` by default (`PeerGet.kind` = `"scope"` / `"all"` switches the view), and are rejected on the generic session-peer routes. Workspace-level key required; peer- and session-scoped keys get 401. Legacy peers occupying a reserved `scope.` name without the kind flag are refused with 409, never adopted (#884)
|
||||
|
|
@ -785,7 +800,7 @@ Welcome to the Honcho changelog! This section documents all notable changes to t
|
|||
|
||||
<Tab title="Python SDK">
|
||||
[Python SDK](https://pypi.org/project/honcho-ai/)
|
||||
<Update label="v2.4.0 (Current)">
|
||||
<Update label="v2.4.0">
|
||||
### Added
|
||||
|
||||
- Scopes: `Honcho.scope()` / `HonchoAio.scope()` get-or-create a named visibility boundary, `Honcho.scopes()` lists them, and a `Scope` object adds/removes sessions, lists membership, and reads backfill `status()`. `Honcho.session(..., scopes=[...])` joins a new session to scopes at creation. Requires a Honcho server with the matching API support (Honcho v3.1.0+).
|
||||
|
|
@ -964,7 +979,7 @@ Welcome to the Honcho changelog! This section documents all notable changes to t
|
|||
|
||||
<Tab title="TypeScript SDK">
|
||||
[TypeScript SDK](https://www.npmjs.com/package/@honcho-ai/sdk)
|
||||
<Update label="v2.4.0 (Current)">
|
||||
<Update label="v2.4.0">
|
||||
### Added
|
||||
|
||||
- Scopes: `honcho.scope()` get-or-creates a named visibility boundary, `honcho.scopes()` lists them, and a `Scope` object adds/removes sessions, lists membership, and reads backfill `status()`. `honcho.session({ scopes: [...] })` joins a new session to scopes at creation. Requires a Honcho server with the matching API support (Honcho v3.1.0+).
|
||||
|
|
@ -1170,7 +1185,7 @@ Welcome to the Honcho changelog! This section documents all notable changes to t
|
|||
</Tab>
|
||||
<Tab title="Honcho CLI">
|
||||
[Honcho CLI](https://pypi.org/project/honcho-cli/)
|
||||
<Update label="v0.1.4 (Current)">
|
||||
<Update label="v0.1.4">
|
||||
### Added
|
||||
|
||||
- A TTY notice when a newer `honcho-cli` is on PyPI (`uv tool upgrade honcho-cli`). Skipped in JSON mode; disable with `HONCHO_NO_UPDATE_CHECK`
|
||||
|
|
|
|||
|
|
@ -28,7 +28,7 @@
|
|||
"navigation": {
|
||||
"versions": [
|
||||
{
|
||||
"version": "v3.1.0",
|
||||
"version": "v3.1.1",
|
||||
"api": {
|
||||
"openapi": ["v3/openapi.json"]
|
||||
},
|
||||
|
|
@ -75,7 +75,8 @@
|
|||
"v3/documentation/features/advanced/using-filters",
|
||||
"v3/documentation/features/advanced/structured-outputs",
|
||||
"v3/documentation/features/advanced/streaming-response",
|
||||
"v3/documentation/features/advanced/file-uploads"
|
||||
"v3/documentation/features/advanced/file-uploads",
|
||||
"v3/documentation/features/advanced/deleting-data"
|
||||
]
|
||||
}
|
||||
]
|
||||
|
|
@ -103,6 +104,7 @@
|
|||
"v3/guides/integrations/claude-code",
|
||||
"v3/guides/integrations/opencode",
|
||||
"v3/guides/integrations/codex",
|
||||
"v3/guides/integrations/deepseek-harness",
|
||||
"v3/guides/integrations/vercel-ai-sdk",
|
||||
"v3/guides/integrations/crewai",
|
||||
"v3/guides/integrations/langgraph",
|
||||
|
|
@ -610,8 +612,8 @@
|
|||
}
|
||||
},
|
||||
"integrations": {
|
||||
"posthog": {
|
||||
"apiKey": "phc_1yrzzcgywqXGcerkkI4g7C0YfyPMcAKNOOvGcjTCiUk"
|
||||
"gtm": {
|
||||
"tagId": "GTM-NSPT9PJF"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,69 @@
|
|||
// Loads PostHog only when the CookieConsent cookie grants Statistics; the
|
||||
// cookie is host-scoped, so a landing-page answer covers the docs.
|
||||
;(function () {
|
||||
var KEY = 'phc_1yrzzcgywqXGcerkkI4g7C0YfyPMcAKNOOvGcjTCiUk'
|
||||
var loaded = false
|
||||
|
||||
function granted() {
|
||||
var m = document.cookie.match(/(?:^|;\s*)CookieConsent=([^;]*)/)
|
||||
if (!m) return false
|
||||
var v = decodeURIComponent(m[1])
|
||||
// "-1" is Cookiebot's consent-not-required marker.
|
||||
return v === '-1' || /statistics\s*:\s*true/.test(v)
|
||||
}
|
||||
|
||||
function loadPosthog() {
|
||||
if (loaded) return
|
||||
loaded = true
|
||||
var s = document.createElement('script')
|
||||
s.src = 'https://us-assets.i.posthog.com/static/array.js'
|
||||
s.async = true
|
||||
s.onerror = function () {
|
||||
loaded = false
|
||||
}
|
||||
s.onload = function () {
|
||||
// Consent withdrawn while array.js was downloading: skip init, allow a retry on re-grant.
|
||||
if (!granted()) {
|
||||
loaded = false
|
||||
return
|
||||
}
|
||||
window.posthog.init(KEY, {
|
||||
api_host: 'https://us.i.posthog.com',
|
||||
ui_host: 'https://us.posthog.com',
|
||||
cross_subdomain_cookie: true,
|
||||
person_profiles: 'identified_only',
|
||||
capture_pageview: 'history_change',
|
||||
})
|
||||
}
|
||||
document.head.appendChild(s)
|
||||
}
|
||||
|
||||
function sync() {
|
||||
if (granted()) {
|
||||
if (!loaded) {
|
||||
loadPosthog()
|
||||
} else if (
|
||||
window.posthog &&
|
||||
window.posthog.has_opted_out_capturing &&
|
||||
window.posthog.has_opted_out_capturing()
|
||||
) {
|
||||
window.posthog.opt_in_capturing()
|
||||
}
|
||||
return
|
||||
}
|
||||
// Withdrawal mid-session: an already running instance must stop.
|
||||
if (loaded && window.posthog && window.posthog.opt_out_capturing) {
|
||||
window.posthog.opt_out_capturing()
|
||||
}
|
||||
}
|
||||
|
||||
sync()
|
||||
var events = [
|
||||
'CookiebotOnConsentReady',
|
||||
'CookiebotOnAccept',
|
||||
'CookiebotOnDecline',
|
||||
]
|
||||
for (var i = 0; i < events.length; i++) {
|
||||
window.addEventListener(events[i], sync)
|
||||
}
|
||||
})()
|
||||
|
|
@ -1,3 +1,13 @@
|
|||
---
|
||||
openapi: post /v3/keys
|
||||
---
|
||||
|
||||
<Note>
|
||||
**Self-hosted only.** This endpoint is not available on Honcho Cloud
|
||||
(`api.honcho.dev`) — requests to it return `405 Method Not Allowed`. Create and
|
||||
manage keys for a cloud instance from the
|
||||
[API Keys page](https://app.honcho.dev/api-keys) in the dashboard.
|
||||
|
||||
On a self-hosted instance it requires an admin key, and returns an error when
|
||||
`AUTH_USE_AUTH` is disabled.
|
||||
</Note>
|
||||
|
|
|
|||
|
|
@ -389,6 +389,20 @@ The default compose file is already production-oriented — ports bound to `127.
|
|||
- You can also run multiple deriver processes across machines — they coordinate via the database queue
|
||||
- Monitor deriver logs for processing backlog
|
||||
|
||||
### Scaling the API
|
||||
|
||||
Set `API_WORKERS` to run multiple API server processes in the Docker container. It defaults to `1`, preserving the existing single-process behavior.
|
||||
|
||||
When connection pooling is enabled (`DB_POOL_CLASS` is not `null`), each API process creates its own SQLAlchemy connection pool. Keep the combined capacity below the PostgreSQL connection limit:
|
||||
|
||||
```text
|
||||
API_WORKERS * (DB_POOL_SIZE + DB_MAX_OVERFLOW) < PostgreSQL max_connections
|
||||
```
|
||||
|
||||
With the default pooled settings (`10 + 20`), each API worker can open up to 30 connections. For example, `API_WORKERS=3` allows up to 90 API connections. Leave additional headroom for the deriver, migrations, administration, and monitoring.
|
||||
|
||||
When `DB_POOL_CLASS=null`, SQLAlchemy uses `NullPool`; `DB_POOL_SIZE` and `DB_MAX_OVERFLOW` do not apply, and connections are opened and closed per use.
|
||||
|
||||
### Caching
|
||||
- The production compose enables Redis caching by default (`CACHE_ENABLED=true`)
|
||||
- For the development compose, enable manually: `CACHE_ENABLED=true`
|
||||
|
|
|
|||
|
|
@ -0,0 +1,131 @@
|
|||
---
|
||||
title: 'Deleting Data'
|
||||
description: 'How to delete sessions, workspaces, and conclusions — and what survives each'
|
||||
icon: 'trash'
|
||||
---
|
||||
|
||||
Deletion in Honcho is **permanent and cannot be undone**. There is no soft
|
||||
delete, no trash, and no restore.
|
||||
|
||||
## What can be deleted
|
||||
|
||||
| Resource | Endpoint | Behavior |
|
||||
|---|---|---|
|
||||
| Session | `DELETE /v3/workspaces/{workspace_id}/sessions/{session_id}` | `202` — cascade runs in the background |
|
||||
| Workspace | `DELETE /v3/workspaces/{workspace_id}` | `202` — cascade runs in the background |
|
||||
| Conclusion | `DELETE /v3/workspaces/{workspace_id}/conclusions/{conclusion_id}` | `204` — immediate |
|
||||
| Webhook endpoint | `DELETE /v3/workspaces/{workspace_id}/webhooks/{endpoint_id}` | Immediate |
|
||||
|
||||
**Peers and individual messages cannot be deleted.** To remove a peer's data,
|
||||
delete the sessions it participated in, then delete its remaining conclusions
|
||||
(see [Conclusions outlive their sessions](#conclusions-outlive-their-sessions)).
|
||||
To remove a peer from one conversation without deleting anything, use
|
||||
[remove peers from session](/v3/api-reference/endpoint/sessions/remove-peers-from-session)
|
||||
instead.
|
||||
|
||||
## Deleting a session
|
||||
|
||||
```bash
|
||||
curl -X DELETE "$HONCHO_URL/v3/workspaces/my-app/sessions/session-1" \
|
||||
-H "Authorization: Bearer $HONCHO_API_KEY"
|
||||
```
|
||||
|
||||
The session is marked inactive immediately and the endpoint returns `202
|
||||
Accepted`. The cascade — messages, message embeddings, queued reasoning work,
|
||||
session-scoped conclusions, and peer associations — is processed in the
|
||||
background with retries.
|
||||
|
||||
Because the work is asynchronous, a `202` means *accepted*, not *finished*. The
|
||||
session drops out of session listings right away, but its messages and
|
||||
conclusions drain afterwards. Deletion tasks are internal infrastructure work
|
||||
and do **not** appear in
|
||||
[queue status](/v3/documentation/features/advanced/queue-status) counts, so
|
||||
there is no endpoint that reports when the cascade has finished.
|
||||
|
||||
<CodeGroup>
|
||||
```python Python
|
||||
session.delete()
|
||||
```
|
||||
|
||||
```typescript TypeScript
|
||||
await session.delete();
|
||||
```
|
||||
</CodeGroup>
|
||||
|
||||
## Deleting a workspace
|
||||
|
||||
A workspace can only be deleted once it has **no active sessions**. Deleting a
|
||||
workspace that still has sessions returns `409 Conflict`:
|
||||
|
||||
```json
|
||||
{"detail": "Cannot delete workspace 'my-app': active session(s) remain. Delete all sessions first."}
|
||||
```
|
||||
|
||||
The correct order is:
|
||||
|
||||
1. List the workspace's sessions — `POST /v3/workspaces/{workspace_id}/sessions/list`
|
||||
2. Delete each session — `DELETE /v3/workspaces/{workspace_id}/sessions/{session_id}`
|
||||
3. Delete the workspace — `DELETE /v3/workspaces/{workspace_id}`
|
||||
|
||||
Step 2 returns `202`, so the session deletions are still draining when step 3
|
||||
runs. That is fine: a session is marked inactive synchronously, so the workspace
|
||||
delete stops returning `409` as soon as the deletes are accepted. Any session
|
||||
created after the workspace deletion is accepted is cascade-deleted too.
|
||||
|
||||
<CodeGroup>
|
||||
```python Python
|
||||
# Materialize the list first — deleting shifts the pagination window
|
||||
for session in list(honcho.sessions()):
|
||||
session.delete()
|
||||
|
||||
honcho.delete_workspace("my-app")
|
||||
```
|
||||
|
||||
```typescript TypeScript
|
||||
// Materialize the list first — deleting shifts the pagination window
|
||||
const sessions = [];
|
||||
for await (const session of await honcho.sessions()) sessions.push(session);
|
||||
for (const session of sessions) await session.delete();
|
||||
|
||||
await honcho.deleteWorkspace("my-app");
|
||||
```
|
||||
</CodeGroup>
|
||||
|
||||
Deleting a workspace removes every peer, session, message, conclusion,
|
||||
collection, embedding, webhook endpoint, and queued task belonging to it.
|
||||
|
||||
## Conclusions outlive their sessions
|
||||
|
||||
This is the most common surprise. Deleting a session does **not** erase
|
||||
everything Honcho learned in it.
|
||||
|
||||
- **Explicit conclusions** — direct facts drawn from messages — are tied to the
|
||||
session they came from and are deleted with it.
|
||||
- **Derived conclusions** (deductive, inductive, contradiction) are consolidations
|
||||
that may draw on several sessions. They are stored at the workspace level with
|
||||
no owning session, so they survive session deletion and stay in the peer's
|
||||
[representation](/v3/documentation/core-concepts/representation).
|
||||
|
||||
To remove those, list and delete them directly:
|
||||
|
||||
<CodeGroup>
|
||||
```python Python
|
||||
for conclusion in alice.conclusions.list():
|
||||
alice.conclusions.delete(conclusion.id)
|
||||
```
|
||||
|
||||
```typescript TypeScript
|
||||
for (const conclusion of await alice.conclusions.list()) {
|
||||
await alice.conclusions.delete(conclusion.id);
|
||||
}
|
||||
```
|
||||
</CodeGroup>
|
||||
|
||||
Deleting the whole workspace removes conclusions at every level and needs no
|
||||
follow-up.
|
||||
|
||||
## Permissions
|
||||
|
||||
Session and workspace deletion accept any key scoped to that workspace — an
|
||||
admin key is not required. Deleting a session additionally accepts a
|
||||
session-scoped key.
|
||||
|
|
@ -23,3 +23,4 @@ Advanced features give you fine-grained control over Honcho's behavior and imple
|
|||
- [Filters](/v3/documentation/features/advanced/using-filters) - Filter queries with advanced parameters
|
||||
- [Streaming Responses](/v3/documentation/features/advanced/streaming-response) - Stream dialectic responses in real-time
|
||||
- [File Uploads](/v3/documentation/features/advanced/file-uploads) - Ingest files into peer memory
|
||||
- [Deleting Data](/v3/documentation/features/advanced/deleting-data) - Delete sessions, workspaces, and conclusions
|
||||
|
|
|
|||
|
|
@ -13,6 +13,14 @@ for a session has drained.
|
|||
Webhooks are registered per workspace. Every event for that workspace is
|
||||
delivered to every endpoint registered on it.
|
||||
|
||||
<Note>
|
||||
**On Honcho Cloud, register endpoints from the dashboard.** The webhook API
|
||||
below is available on self-hosted instances; on `api.honcho.dev` it returns
|
||||
`405 Method Not Allowed`. Use the
|
||||
[Webhooks page](https://app.honcho.dev/webhooks) instead. Everything else on
|
||||
this page — payload shapes, delivery semantics — applies to both.
|
||||
</Note>
|
||||
|
||||
## Registering an Endpoint
|
||||
|
||||
<CodeGroup>
|
||||
|
|
|
|||
|
|
@ -62,6 +62,8 @@ The **Performance** page provides comprehensive monitoring with usage metrics, h
|
|||
## 3. Manage API Keys
|
||||
The [API Keys](https://app.honcho.dev/api-keys) page allows you to create and manage authentication tokens for different environments. You can create admin-level keys with full instance access or scope keys to a specific `Workspace`, `Peer`, or `Session`.
|
||||
|
||||
Keys for a cloud instance can only be created here, not through the API — `POST /v3/keys` is disabled on `api.honcho.dev` and returns `405`. The same applies to the webhook management endpoints, which live on the [Webhooks](https://app.honcho.dev/webhooks) page.
|
||||
|
||||
Scoped keys are authorized by their narrowest claim and never widen to the whole workspace:
|
||||
|
||||
- A **peer-scoped** key acts on its own peer, plus **read-only** access to the sessions its peer is an active member of (context, summaries, peers, its own per-session config, search, and message reads). It cannot write to those sessions or act on other peers.
|
||||
|
|
|
|||
|
|
@ -206,6 +206,9 @@ honcho.set_metadata(dict)
|
|||
|
||||
# Get list of all workspace IDs
|
||||
workspaces = honcho.workspaces()
|
||||
|
||||
# Delete a workspace and everything in it (requires no active sessions)
|
||||
honcho.delete_workspace(workspace_id)
|
||||
```
|
||||
|
||||
```typescript TypeScript
|
||||
|
|
@ -238,6 +241,9 @@ await honcho.setMetadata(metadata);
|
|||
|
||||
// Get list of all workspace IDs
|
||||
const workspaces = await honcho.workspaces();
|
||||
|
||||
// Delete a workspace and everything in it (requires no active sessions)
|
||||
await honcho.deleteWorkspace(workspaceId);
|
||||
```
|
||||
</CodeGroup>
|
||||
|
||||
|
|
@ -270,6 +276,9 @@ response = alice.chat("What do I know about Bob?", target="bob")
|
|||
response = alice.chat("What happened in session-1?", session="session-1")
|
||||
response = alice.chat("Summarize what matters most to me.", reasoning_level="high")
|
||||
|
||||
# Override the timeout for one non-streaming dialectic request
|
||||
response = alice.chat("Give me a quick summary.", timeout=5.0)
|
||||
|
||||
# Add content to a session with a peer
|
||||
session = honcho.session("session-1")
|
||||
session.add_messages([
|
||||
|
|
@ -372,6 +381,13 @@ const bobConclusions = await alice.conclusionsOf("bob").list(); // Conclusions
|
|||
```
|
||||
</CodeGroup>
|
||||
|
||||
For Python, `peer.chat(timeout=...)` and `await peer.aio.chat(timeout=...)`
|
||||
accept a timeout in seconds for each HTTP attempt made by one non-streaming
|
||||
request. Omit it or pass `None` to use the client-wide timeout configured on
|
||||
`Honcho`. Retries still follow the client's `max_retries` setting and can extend
|
||||
total elapsed time; use `max_retries=0` when a host shutdown budget permits only
|
||||
one attempt.
|
||||
|
||||
### Peer Context
|
||||
|
||||
The `context()` method on peers retrieves both the working representation and peer card in a single API call:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,198 @@
|
|||
---
|
||||
title: "DeepSeek Harness"
|
||||
icon: 'terminal'
|
||||
description: "Add AI-native memory to DeepSeek Harness"
|
||||
sidebarTitle: 'DeepSeek Harness'
|
||||
---
|
||||
|
||||
`dsh` forgets everything when a session ends. This plugin gives it memory that doesn't: what you're building, how you like to work, and what you decided last week and why — carried across context resets, restarts, and fresh chats.
|
||||
|
||||
It is a native [Cordis](https://github.com/cordiverse/cordis) plugin, not a hook bridge, so it hooks the harness's own extension points directly.
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Step 1: Get Your Honcho API Key
|
||||
|
||||
1. Go to **[app.honcho.dev](https://app.honcho.dev)**
|
||||
2. Sign up or log in
|
||||
3. Copy your API key (starts with `hch-`)
|
||||
|
||||
### Step 2: Install the Plugin
|
||||
|
||||
<Note>
|
||||
This plugin requires a running [DeepSeek Harness](https://github.com/deepseek-ai/deepseek-harness). Plugins install into a named profile, so pick the one you actually run — `web`, `headless`, `acp`, or your own.
|
||||
</Note>
|
||||
|
||||
```bash
|
||||
dsh plugin --profile <name> add @honcho-ai/dsh-honcho
|
||||
```
|
||||
|
||||
`dsh plugin` forwards to your package manager and appends the plugin to that profile's bundle list. Because the package declares `dsh.bundle`, it activates as a configuration layer rather than sitting inert as a plain dependency.
|
||||
|
||||
### Step 3: Configure
|
||||
|
||||
Put your key and name in `~/.honcho/config.json`:
|
||||
|
||||
```jsonc
|
||||
{
|
||||
"peerName": "your-name",
|
||||
"auth": { "apiKey": "${HONCHO_API_KEY}" },
|
||||
"hosts": {
|
||||
"dsh": { "workspace": "dsh" }
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
`HONCHO_API_KEY` in the environment works on its own — the config file is only needed to change defaults.
|
||||
|
||||
### Step 4: Verify
|
||||
|
||||
Start `dsh` and run `/honcho`. You'll see your peer, workspace, session, and sync status, plus a link to the session in the Honcho dashboard.
|
||||
|
||||
<Warning>
|
||||
In the `dsh` web client, `/honcho` output renders in the collapsed command panel rather than inline in the transcript. Expand the panel to read it.
|
||||
</Warning>
|
||||
|
||||
## What You Get
|
||||
|
||||
- **Memory at session start** — your profile, a summary of this project's session so far, and the conclusions relevant to what you just asked, shaped to a character budget in a single API call
|
||||
- **Automatic capture** — user and assistant turns stream to Honcho in the background, debounced, and flushed at turn boundaries, before compaction, and on shutdown
|
||||
- **Secret redaction** — messages are scrubbed before they leave your machine
|
||||
- **Agent tools** — first-class search, reasoning, and conclusion-writing inside `dsh`
|
||||
- **Shared configuration** — the same `~/.honcho/config.json` every other Honcho integration reads
|
||||
|
||||
## Configuration
|
||||
|
||||
Configuration lives in `~/.honcho/config.json`, shared with the other Honcho hosts. The root holds identity and connection; behavior lives under `hosts.dsh`.
|
||||
|
||||
```jsonc
|
||||
{
|
||||
"peerName": "your-name",
|
||||
"workspace": "honcho",
|
||||
"baseUrl": "https://api.honcho.dev", // bare host or …/v3 both fine
|
||||
"timeoutMs": 30000,
|
||||
"auth": { "apiKey": "${HONCHO_API_KEY}" },
|
||||
"enabled": true, // global kill switch
|
||||
|
||||
"hosts": {
|
||||
"dsh": {
|
||||
"workspace": "dsh",
|
||||
"aiPeer": "dsh", // defaults to the host name
|
||||
"observationMode": "unified", // unified | directional
|
||||
"sessionStrategy": "per-directory",
|
||||
"sessionPeerPrefix": true, // session names are <peer>-<dir>
|
||||
"sessions": { "/path/to/repo": "pinned-session-name" },
|
||||
|
||||
"injection": {
|
||||
"sessionStart": ["directives", "summary", "peerCard"],
|
||||
"perTurn": ["userContext", "dialectic"],
|
||||
"tools": true,
|
||||
"searchTopK": 10,
|
||||
"searchMaxDistance": 0.6,
|
||||
"maxConclusions": 15, // how many conclusions Honcho RETURNS
|
||||
"maxRenderedConclusions": 4, // how many survive into the prompt
|
||||
"contextTokens": 1500,
|
||||
"cadence": { "dialectic": 5, "ttlSeconds": 300 },
|
||||
"dialectic": {
|
||||
"reasoning": "low", // minimal | low | medium | high | max
|
||||
"maxChars": 600
|
||||
}
|
||||
},
|
||||
|
||||
"capture": {
|
||||
"saveMessages": true,
|
||||
"saveToolUse": false, // one-line summaries of tool activity
|
||||
"writeFrequency": "async", // async | sync
|
||||
"noisePatterns": [] // additive to the built-in secret patterns
|
||||
},
|
||||
|
||||
"messageUpload": {
|
||||
"maxUserTokens": 6000,
|
||||
"maxAssistantTokens": 6000
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
<Note>
|
||||
Unsupported or renamed keys are reported at startup rather than silently ignored, so a stale config tells you what it is no longer doing.
|
||||
</Note>
|
||||
|
||||
### Injection Components
|
||||
|
||||
The two menus differ in **cadence**, not in what they can carry.
|
||||
|
||||
`injection.sessionStart` is injected once when a session opens: `directives`, `summary`, `peerCard`, `representation`.
|
||||
|
||||
`injection.perTurn` refreshes as you work:
|
||||
|
||||
| Component | Behavior |
|
||||
| --- | --- |
|
||||
| `userContext` | A fresh, prompt-scoped bundle of **representation + peer card**, retrieved using your current message as the search query — so recall is associative rather than merely recent |
|
||||
| `dialectic` | A reasoned answer about you, run every `cadence.dialectic` turns. Nothing waits on it after the first turn, so a late answer reaches the next one |
|
||||
|
||||
To get the representation without the peer card (or vice versa), name it in `sessionStart` and set `perTurn: []` — at the cost of per-turn refresh.
|
||||
|
||||
### Session Strategies
|
||||
|
||||
| Strategy | Session name | Notes |
|
||||
| --- | --- | --- |
|
||||
| `per-directory` (default) | `<peer>-<dir>` | Stable across restarts and branches |
|
||||
| `per-repo` | `<peer>-<repo-root>` | Same memory from any subdirectory |
|
||||
| `git-branch` | `<peer>-<dir>-<branch>` | Falls back to `per-directory` outside a repo or on a detached HEAD |
|
||||
| `per-session` | `<peer>-chat-<id>` | A clean slate every restart |
|
||||
| `global` | `<peer>` | One memory for everything |
|
||||
|
||||
<Warning>
|
||||
Prefer the wider scopes. The background Deriver needs a single session to accumulate enough material before it can reason well. `git-branch` splits a project's memory per branch, and `per-session` discards it on every restart.
|
||||
</Warning>
|
||||
|
||||
### Sharing Memory With Other Integrations
|
||||
|
||||
Each integration defaults to its own Honcho `workspace` — `dsh` here, `claude_code` for claude-honcho — and a workspace is the isolation boundary, so **by default they do not see each other's memory.** Point them at the same `workspace` to merge them:
|
||||
|
||||
```jsonc
|
||||
"hosts": {
|
||||
"dsh": { "workspace": "shared" },
|
||||
"claude_code": { "workspace": "shared" }
|
||||
}
|
||||
```
|
||||
|
||||
Keep `peerName` identical across them too, since conclusions are stored per peer.
|
||||
|
||||
## Commands
|
||||
|
||||
| Command | Description |
|
||||
| --- | --- |
|
||||
| `/honcho` | Status: peer, workspace, session, strategy, pending uploads, last sync, last fetch |
|
||||
| `/honcho config` | Resolved settings, the file they came from, and any ignored injection components |
|
||||
| `/honcho flush` | Sync now |
|
||||
|
||||
## Agent Tools
|
||||
|
||||
| Tool | Description |
|
||||
| --- | --- |
|
||||
| `honcho_search` | Look something up — searches raw messages **and** derived conclusions |
|
||||
| `honcho_chat` | Ask a question of judgment. Reasons over everything Honcho knows; slower |
|
||||
| `honcho_remember` | Save a durable fact, preference, or decision |
|
||||
|
||||
Set `injection.tools` to `false` to inject memory without exposing tools.
|
||||
|
||||
## Requirements
|
||||
|
||||
- Node `^22.19.0 || >=24.0.0`
|
||||
- A running `dsh`
|
||||
- A Honcho API key, or a self-hosted Honcho at `baseUrl`
|
||||
|
||||
## Next Steps
|
||||
|
||||
<CardGroup cols={2}>
|
||||
<Card title="GitHub Repository" icon="github" href="https://github.com/plastic-labs/dsh-honcho">
|
||||
Source code, issues, and README.
|
||||
</Card>
|
||||
|
||||
<Card title="Honcho Architecture" icon="sitemap" href="../../documentation/core-concepts/architecture">
|
||||
Learn about peers, sessions, and dialectic reasoning.
|
||||
</Card>
|
||||
</CardGroup>
|
||||
|
|
@ -9,7 +9,7 @@
|
|||
"url": "https://honcho.dev/",
|
||||
"email": "hello@plasticlabs.ai"
|
||||
},
|
||||
"version": "3.1.0"
|
||||
"version": "3.1.1"
|
||||
},
|
||||
"servers": [
|
||||
{
|
||||
|
|
|
|||
|
|
@ -0,0 +1,2 @@
|
|||
node_modules
|
||||
dist
|
||||
|
|
@ -0,0 +1,10 @@
|
|||
# Changelog
|
||||
|
||||
All notable changes to `@honcho-ai/harness-plugin-core` will be documented in this file.
|
||||
|
||||
The format is based on [Keep a Changelog](http://keepachangelog.com/)
|
||||
and this project adheres to [Semantic Versioning](http://semver.org/).
|
||||
|
||||
This package versions independently of the Honcho API, `@honcho-ai/sdk`, and host plugins.
|
||||
|
||||
## [Unreleased]
|
||||
|
|
@ -0,0 +1,70 @@
|
|||
# @honcho-ai/harness-plugin-core
|
||||
|
||||
Shared runtime for Honcho harness plugins.
|
||||
|
||||
```ts
|
||||
import { loadConfig, resolveConfig } from '@honcho-ai/harness-plugin-core'
|
||||
|
||||
const cfg = loadConfig({ host: 'harness' })
|
||||
// a harness can pass its plugin config as an overlay of the same six keys:
|
||||
const cfg = resolveConfig(file, { host: 'harness', overlay: { workspace: 'harness', auth: { apiKey } } })
|
||||
```
|
||||
|
||||
Locally: `"@honcho-ai/harness-plugin-core": "file:../harness-plugin-core"` (bun imports the TypeScript source).
|
||||
|
||||
## File shape
|
||||
|
||||
```json
|
||||
{
|
||||
"schemaVersion": 1,
|
||||
"peerName": "user",
|
||||
"workspace": "honcho",
|
||||
"baseUrl": "https://api.honcho.dev",
|
||||
"timeoutMs": 30000,
|
||||
"auth": { "apiKey": "${HONCHO_API_KEY}" },
|
||||
"enabled": true,
|
||||
"hosts": {
|
||||
"test": { "workspace": "test" }
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Missing `schemaVersion` is 0. On read, v0 keys (`environmentUrl`, `workspaceId`, top-level `apiKey`) are remapped in memory; the file is not rewritten.
|
||||
|
||||
Resolution, highest wins: `HONCHO_*` env → overlay → `hosts.<host>` → root → built-in.
|
||||
|
||||
A host block may override the same six fields.
|
||||
|
||||
Built-ins: `baseUrl = https://api.honcho.dev`, `timeoutMs = 30000`, `enabled = true`, `peerName = $USER`, `workspace` falls back to the host name. The SDK pins `/v3`; config stores the origin.
|
||||
|
||||
## Telemetry headers
|
||||
|
||||
Pass `telemetryHeaders()` as the SDK's `defaultHeaders`. Arbitrary headers are accepted by both the SDK and the Honcho API; missing identity fields are omitted.
|
||||
|
||||
| Header | Meaning | Example |
|
||||
|---|---|---|
|
||||
| `X-Honcho-Host` | Agent host name, or `name/version` | `harness/1.3.13` |
|
||||
| `X-Honcho-Plugin` | Honcho plugin version | `0.1.3` |
|
||||
| `X-Honcho-Runtime` | This package's version (always sent) | `0.1.0` |
|
||||
| `X-Honcho-Agent-Model` | The agent's completion model, not a Honcho model | `claude-sonnet-4-5` |
|
||||
|
||||
```ts
|
||||
import { Honcho } from '@honcho-ai/sdk'
|
||||
import { loadConfig, setTelemetryHeaders, telemetryHeaders } from '@honcho-ai/harness-plugin-core'
|
||||
|
||||
const cfg = loadConfig({ host: 'harness' })
|
||||
const honcho = new Honcho({
|
||||
apiKey: cfg.apiKey,
|
||||
baseURL: cfg.baseUrl,
|
||||
workspaceId: cfg.workspace,
|
||||
timeout: cfg.timeoutMs,
|
||||
defaultHeaders: telemetryHeaders({
|
||||
host: 'harness',
|
||||
hostVersion: '1.3.13',
|
||||
pluginVersion: '0.1.3',
|
||||
model: 'claude-sonnet-4-5',
|
||||
}),
|
||||
})
|
||||
|
||||
setTelemetryHeaders(honcho.http.defaultHeaders, { model: 'claude-opus-4' })
|
||||
```
|
||||
|
|
@ -0,0 +1,25 @@
|
|||
{
|
||||
"lockfileVersion": 1,
|
||||
"configVersion": 1,
|
||||
"workspaces": {
|
||||
"": {
|
||||
"name": "@honcho-ai/harness-plugin-core",
|
||||
"devDependencies": {
|
||||
"@types/bun": "latest",
|
||||
"@types/node": "^24.0.1",
|
||||
"typescript": "^5.0.0",
|
||||
},
|
||||
},
|
||||
},
|
||||
"packages": {
|
||||
"@types/bun": ["@types/bun@1.4.0", "", { "dependencies": { "bun-types": "1.4.0" } }, "sha512-K+lZULY23vRgK/CfTjFIV+tyifaNdSMlPh9j+6mQ/cLfpOznLyAuzgV/JQysyECpkBQLVMSyvjlr2fBUSA9wFQ=="],
|
||||
|
||||
"@types/node": ["@types/node@24.13.3", "", { "dependencies": { "undici-types": "~7.18.0" } }, "sha512-Dh8vAsV36ig5wa9OX4pXvMc9D3Veibfw2wix0CUwYODLD8nkj9UsLjASr49nPg+2eKzxhBV+v7L8pXvT4e639Q=="],
|
||||
|
||||
"bun-types": ["bun-types@1.4.0", "", { "dependencies": { "@types/node": "*" } }, "sha512-iIKw23BspnQQYd3prITOBxeUsxBHnwzX6YJfGMuNOZzeNcMmVqzIIVGRm1l69ogaPQmb4wB6BN8mA5bE9YuC5Q=="],
|
||||
|
||||
"typescript": ["typescript@5.9.3", "", { "bin": { "tsc": "bin/tsc", "tsserver": "bin/tsserver" } }, "sha512-jl1vZzPDinLr9eUt3J/t7V6FgNEw9QjvBPdysz9KfQDD41fQrC2Y4vKQdiaUpFT4bXlb1RHhLpp8wtm6M5TgSw=="],
|
||||
|
||||
"undici-types": ["undici-types@7.18.2", "", {}, "sha512-AsuCzffGHJybSaRrmr5eHr81mwJU3kjw6M+uprWvCXiNeN9SOGwQ3Jn8jb8m3Z6izVgknn1R0FTCEAP2QrLY/w=="],
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,31 @@
|
|||
{
|
||||
"name": "@honcho-ai/harness-plugin-core",
|
||||
"version": "0.1.0",
|
||||
"description": "Shared runtime for Honcho harness plugins",
|
||||
"author": "Plastic Labs <hello@plasticlabs.ai>",
|
||||
"license": "MIT",
|
||||
"type": "module",
|
||||
"main": "src/index.ts",
|
||||
"exports": {
|
||||
".": "./src/index.ts"
|
||||
},
|
||||
"files": [
|
||||
"src",
|
||||
"README.md",
|
||||
"CHANGELOG.md"
|
||||
],
|
||||
"repository": {
|
||||
"type": "git",
|
||||
"url": "git+https://github.com/plastic-labs/honcho.git",
|
||||
"directory": "harness-plugin-core"
|
||||
},
|
||||
"scripts": {
|
||||
"test": "bun test",
|
||||
"typecheck": "tsc --noEmit"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@types/bun": "latest",
|
||||
"@types/node": "^24.0.1",
|
||||
"typescript": "^5.0.0"
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,239 @@
|
|||
import { existsSync, readFileSync } from 'node:fs'
|
||||
import { homedir } from 'node:os'
|
||||
import { join } from 'node:path'
|
||||
|
||||
export interface AuthConfig {
|
||||
apiKey?: string
|
||||
oauth?: { accessToken?: string; refreshToken?: string; expiresAt?: string }
|
||||
}
|
||||
|
||||
/** Identity + connection + kill switch. Valid at root and as a host override. */
|
||||
export interface RootConfig {
|
||||
peerName?: string
|
||||
workspace?: string
|
||||
baseUrl?: string
|
||||
timeoutMs?: number
|
||||
auth?: AuthConfig
|
||||
enabled?: boolean
|
||||
}
|
||||
|
||||
export type HostBlock = RootConfig
|
||||
|
||||
export interface FileConfig extends RootConfig {
|
||||
schemaVersion?: number
|
||||
hosts?: Record<string, HostBlock>
|
||||
}
|
||||
|
||||
export interface ResolvedConfig {
|
||||
host: string
|
||||
peerName: string
|
||||
workspace: string
|
||||
baseUrl: string
|
||||
timeoutMs: number
|
||||
auth: AuthConfig
|
||||
apiKey?: string
|
||||
enabled: boolean
|
||||
warnings: string[]
|
||||
}
|
||||
|
||||
export const DEFAULT_BASE_URL = 'https://api.honcho.dev'
|
||||
export const DEFAULT_TIMEOUT_MS = 30_000
|
||||
export const CONFIG_SCHEMA_VERSION = 1
|
||||
|
||||
function isObj(v: unknown): v is Record<string, unknown> {
|
||||
return v !== null && typeof v === 'object' && !Array.isArray(v)
|
||||
}
|
||||
|
||||
/** Pre-schema files (no schemaVersion) → v1 keys. Host blocks included. */
|
||||
function migrate(file: unknown): Record<string, unknown> {
|
||||
if (!isObj(file)) return {}
|
||||
const v = file.schemaVersion
|
||||
if (typeof v === 'number' && v >= CONFIG_SCHEMA_VERSION) return { ...file }
|
||||
const out: Record<string, unknown> = { ...file }
|
||||
const blocks: Record<string, unknown>[] = [out]
|
||||
if (isObj(out.hosts)) {
|
||||
out.hosts = Object.fromEntries(
|
||||
Object.entries(out.hosts).map(([k, block]) => {
|
||||
if (!isObj(block)) return [k, block]
|
||||
const next = { ...block }
|
||||
blocks.push(next)
|
||||
return [k, next]
|
||||
})
|
||||
)
|
||||
}
|
||||
for (const b of blocks) {
|
||||
if (typeof b.baseUrl !== 'string') {
|
||||
if (typeof b.environmentUrl === 'string') b.baseUrl = b.environmentUrl
|
||||
else if (isObj(b.endpoint) && typeof b.endpoint.baseUrl === 'string') {
|
||||
b.baseUrl = b.endpoint.baseUrl
|
||||
}
|
||||
}
|
||||
if (typeof b.workspace !== 'string' && typeof b.workspaceId === 'string') {
|
||||
b.workspace = b.workspaceId
|
||||
}
|
||||
const auth: Record<string, unknown> = isObj(b.auth) ? { ...b.auth } : {}
|
||||
if (typeof auth.apiKey !== 'string' && typeof b.apiKey === 'string') auth.apiKey = b.apiKey
|
||||
if (!isObj(auth.oauth) && isObj(b.oauth)) auth.oauth = b.oauth
|
||||
if (Object.keys(auth).length) b.auth = auth
|
||||
delete b.environmentUrl
|
||||
delete b.endpoint
|
||||
delete b.workspaceId
|
||||
delete b.apiKey
|
||||
delete b.oauth
|
||||
}
|
||||
out.schemaVersion = 1
|
||||
return out
|
||||
}
|
||||
|
||||
function merge<T>(base: T, over: unknown): T {
|
||||
if (over === undefined || over === null) return base
|
||||
if (Array.isArray(over) || !isObj(over)) return over as T
|
||||
const out: Record<string, unknown> = { ...(isObj(base) ? base : {}) }
|
||||
for (const [k, v] of Object.entries(over)) {
|
||||
if (v !== undefined) out[k] = k in out ? merge(out[k], v) : v
|
||||
}
|
||||
return out as T
|
||||
}
|
||||
|
||||
/** Make a value safe to pass to the SDK as `baseURL`. */
|
||||
export function normalizeBaseUrl(input: string): string {
|
||||
let s = input.trim()
|
||||
if (!s) return s
|
||||
if (!s.startsWith('http://') && !s.startsWith('https://')) {
|
||||
const host = s.split('/')[0].split(':')[0].toLowerCase()
|
||||
const local = host === 'localhost' || host === '127.0.0.1' || host === '::1'
|
||||
s = `${local ? 'http' : 'https'}://${s}`
|
||||
}
|
||||
try {
|
||||
const u = new URL(s)
|
||||
u.hostname = u.hostname.toLowerCase()
|
||||
const path = u.pathname === '/' ? '' : u.pathname.replace(/\/+$/, '')
|
||||
return `${u.protocol}//${u.host}${path}`
|
||||
} catch {
|
||||
return s
|
||||
}
|
||||
}
|
||||
|
||||
function interpolate(value: string, env: NodeJS.Dict<string>, warnings: string[]): string {
|
||||
return value.replace(/\$\{([^}]+)\}/g, (m, name: string) => {
|
||||
const v = env[name]
|
||||
if (!v) {
|
||||
warnings.push(`${m} is not set`)
|
||||
return m
|
||||
}
|
||||
return v
|
||||
})
|
||||
}
|
||||
|
||||
function walkStrings<T>(value: T, fn: (s: string) => string): T {
|
||||
if (typeof value === 'string') return fn(value) as T
|
||||
if (Array.isArray(value)) return value.map((x) => walkStrings(x, fn)) as T
|
||||
if (isObj(value)) {
|
||||
const out: Record<string, unknown> = {}
|
||||
for (const [k, v] of Object.entries(value)) out[k] = walkStrings(v, fn)
|
||||
return out as T
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
/** Pull only the six root fields. Extra host keys (injection, observation, …) are ignored. */
|
||||
function pickRoot(block: unknown): RootConfig {
|
||||
if (!isObj(block)) return {}
|
||||
const auth: AuthConfig = isObj(block.auth) ? { ...(block.auth as AuthConfig) } : {}
|
||||
const out: RootConfig = {}
|
||||
if (typeof block.peerName === 'string') out.peerName = block.peerName
|
||||
if (typeof block.workspace === 'string') out.workspace = block.workspace
|
||||
if (typeof block.baseUrl === 'string') out.baseUrl = block.baseUrl
|
||||
if (typeof block.timeoutMs === 'number') out.timeoutMs = block.timeoutMs
|
||||
if (Object.keys(auth).length) out.auth = auth
|
||||
if (typeof block.enabled === 'boolean') out.enabled = block.enabled
|
||||
return out
|
||||
}
|
||||
|
||||
function pickHost(hosts: Record<string, unknown> | undefined, name: string): RootConfig {
|
||||
if (!hosts || !isObj(hosts[name])) return {}
|
||||
return pickRoot(hosts[name])
|
||||
}
|
||||
|
||||
/**
|
||||
* Highest wins: HONCHO_* env → overlay → hosts.<host> → root → built-in.
|
||||
*/
|
||||
export function resolveConfig(
|
||||
file: unknown,
|
||||
opts: { host: string; env?: NodeJS.Dict<string>; overlay?: RootConfig }
|
||||
): ResolvedConfig {
|
||||
const warnings: string[] = []
|
||||
const env = opts.env ?? process.env
|
||||
const host = opts.host
|
||||
const raw = migrate(file)
|
||||
if (typeof raw.schemaVersion === 'number' && raw.schemaVersion > CONFIG_SCHEMA_VERSION) {
|
||||
warnings.push(`config schemaVersion ${raw.schemaVersion} is newer than ${CONFIG_SCHEMA_VERSION}`)
|
||||
}
|
||||
const hosts = isObj(raw.hosts) ? raw.hosts : undefined
|
||||
|
||||
let acc: RootConfig = {
|
||||
baseUrl: DEFAULT_BASE_URL,
|
||||
timeoutMs: DEFAULT_TIMEOUT_MS,
|
||||
enabled: true,
|
||||
workspace: host,
|
||||
}
|
||||
acc = merge(acc, pickRoot(raw))
|
||||
acc = merge(acc, pickHost(hosts, host))
|
||||
acc = merge(acc, pickRoot(opts.overlay))
|
||||
|
||||
if (env.HONCHO_API_KEY) {
|
||||
if (acc.auth?.apiKey) warnings.push('HONCHO_API_KEY shadows auth.apiKey')
|
||||
acc = merge(acc, { auth: { apiKey: env.HONCHO_API_KEY } })
|
||||
}
|
||||
if (env.HONCHO_BASE_URL || env.HONCHO_URL || env.HONCHO_ENDPOINT) {
|
||||
const token = env.HONCHO_BASE_URL || env.HONCHO_URL || env.HONCHO_ENDPOINT || ''
|
||||
acc.baseUrl = token === 'local' ? 'http://localhost:8000' : token
|
||||
}
|
||||
if (env.HONCHO_WORKSPACE || env.HONCHO_WORKSPACE_ID) {
|
||||
acc.workspace = env.HONCHO_WORKSPACE || env.HONCHO_WORKSPACE_ID
|
||||
}
|
||||
if (env.HONCHO_PEER_NAME) acc.peerName = env.HONCHO_PEER_NAME
|
||||
if (env.HONCHO_TIMEOUT_MS) {
|
||||
const n = Number(env.HONCHO_TIMEOUT_MS)
|
||||
if (Number.isFinite(n) && n > 0) acc.timeoutMs = n
|
||||
}
|
||||
if (env.HONCHO_ENABLED === 'false') acc.enabled = false
|
||||
|
||||
acc = walkStrings(acc, (s) => interpolate(s, env, warnings))
|
||||
if (acc.baseUrl) acc.baseUrl = normalizeBaseUrl(acc.baseUrl)
|
||||
|
||||
const auth = acc.auth ?? {}
|
||||
return {
|
||||
host,
|
||||
peerName: acc.peerName || env.USER || env.USERNAME || 'user',
|
||||
workspace: acc.workspace || host,
|
||||
baseUrl: acc.baseUrl || DEFAULT_BASE_URL,
|
||||
timeoutMs: acc.timeoutMs && acc.timeoutMs > 0 ? acc.timeoutMs : DEFAULT_TIMEOUT_MS,
|
||||
auth,
|
||||
apiKey: auth.apiKey,
|
||||
enabled: acc.enabled !== false,
|
||||
warnings,
|
||||
}
|
||||
}
|
||||
|
||||
export function configPath(env: NodeJS.Dict<string> = process.env): string {
|
||||
return env.HONCHO_CONFIG_PATH || join(homedir(), '.honcho', 'config.json')
|
||||
}
|
||||
|
||||
export function loadConfig(opts: {
|
||||
host: string
|
||||
env?: NodeJS.Dict<string>
|
||||
overlay?: RootConfig
|
||||
}): ResolvedConfig {
|
||||
const env = opts.env ?? process.env
|
||||
const path = configPath(env)
|
||||
let file: unknown = {}
|
||||
if (existsSync(path)) {
|
||||
try {
|
||||
file = JSON.parse(readFileSync(path, 'utf-8'))
|
||||
} catch {
|
||||
file = {}
|
||||
}
|
||||
}
|
||||
return resolveConfig(file, { ...opts, env })
|
||||
}
|
||||
|
|
@ -0,0 +1,29 @@
|
|||
export const version = '0.1.0'
|
||||
|
||||
export {
|
||||
configPath,
|
||||
loadConfig,
|
||||
normalizeBaseUrl,
|
||||
resolveConfig,
|
||||
DEFAULT_BASE_URL,
|
||||
DEFAULT_TIMEOUT_MS,
|
||||
} from './config.ts'
|
||||
|
||||
export type {
|
||||
AuthConfig,
|
||||
FileConfig,
|
||||
HostBlock,
|
||||
ResolvedConfig,
|
||||
RootConfig,
|
||||
} from './config.ts'
|
||||
|
||||
export {
|
||||
telemetryHeaders,
|
||||
setTelemetryHeaders,
|
||||
HEADER_AGENT_MODEL,
|
||||
HEADER_HOST,
|
||||
HEADER_PLUGIN,
|
||||
HEADER_RUNTIME,
|
||||
} from './telemetry.ts'
|
||||
|
||||
export type { TelemetryIdentity } from './telemetry.ts'
|
||||
|
|
@ -0,0 +1,64 @@
|
|||
import { version } from './index.ts'
|
||||
|
||||
/** Optional identity a host plugin knows at Honcho-client construction time. */
|
||||
export interface TelemetryIdentity {
|
||||
/** Host app name, e.g. `cursor`, `opencode`. */
|
||||
host?: string
|
||||
/** Host app version, e.g. `2026.8.1`. */
|
||||
hostVersion?: string
|
||||
/** Honcho plugin version, e.g. `0.1.2`. */
|
||||
pluginVersion?: string
|
||||
/** Agent completion model, e.g. `claude-sonnet-4-5`. Not a Honcho deriver/dialectic model. */
|
||||
model?: string
|
||||
}
|
||||
|
||||
export const HEADER_HOST = 'X-Honcho-Host'
|
||||
export const HEADER_PLUGIN = 'X-Honcho-Plugin'
|
||||
export const HEADER_RUNTIME = 'X-Honcho-Runtime'
|
||||
export const HEADER_AGENT_MODEL = 'X-Honcho-Agent-Model'
|
||||
|
||||
function sanitize(value: unknown): string | undefined {
|
||||
if (typeof value !== 'string') return undefined
|
||||
const s = value.replace(/[\r\n]+/g, ' ').trim()
|
||||
return s || undefined
|
||||
}
|
||||
|
||||
function hostValue(id: TelemetryIdentity): string | undefined {
|
||||
const name = sanitize(id.host)
|
||||
const ver = sanitize(id.hostVersion)
|
||||
if (name && ver) return `${name}/${ver}`
|
||||
return name || ver
|
||||
}
|
||||
|
||||
/**
|
||||
* Headers to pass as the SDK's `defaultHeaders`. Missing fields are omitted.
|
||||
* `X-Honcho-Runtime` is always this package's version.
|
||||
*/
|
||||
export function telemetryHeaders(
|
||||
id: TelemetryIdentity = {},
|
||||
extra?: Record<string, string>
|
||||
): Record<string, string> {
|
||||
const headers: Record<string, string> = { [HEADER_RUNTIME]: version }
|
||||
const host = hostValue(id)
|
||||
const plugin = sanitize(id.pluginVersion)
|
||||
const model = sanitize(id.model)
|
||||
if (host) headers[HEADER_HOST] = host
|
||||
if (plugin) headers[HEADER_PLUGIN] = plugin
|
||||
if (model) headers[HEADER_AGENT_MODEL] = model
|
||||
if (extra) {
|
||||
for (const [k, v] of Object.entries(extra)) {
|
||||
const value = sanitize(v)
|
||||
if (value) headers[k] = value
|
||||
}
|
||||
}
|
||||
return headers
|
||||
}
|
||||
|
||||
/** Merge identity onto a live header map (e.g. `honcho.http.defaultHeaders`). */
|
||||
export function setTelemetryHeaders(
|
||||
headers: Record<string, string>,
|
||||
id: TelemetryIdentity = {},
|
||||
extra?: Record<string, string>
|
||||
): Record<string, string> {
|
||||
return Object.assign(headers, telemetryHeaders(id, extra))
|
||||
}
|
||||
|
|
@ -0,0 +1,71 @@
|
|||
import { describe, expect, test } from 'bun:test'
|
||||
import { normalizeBaseUrl, resolveConfig } from '../src/index.ts'
|
||||
|
||||
const emptyEnv = {}
|
||||
|
||||
describe('normalizeBaseUrl', () => {
|
||||
test('adds https and lowercases the host', () => {
|
||||
expect(normalizeBaseUrl('api.honcho.dev')).toBe('https://api.honcho.dev')
|
||||
expect(normalizeBaseUrl('API.honcho.dev')).toBe('https://api.honcho.dev')
|
||||
expect(normalizeBaseUrl('https://api.honcho.dev/')).toBe('https://api.honcho.dev')
|
||||
})
|
||||
|
||||
test('leaves /v3 alone — the SDK owns the API version', () => {
|
||||
expect(normalizeBaseUrl('https://api.honcho.dev/v3')).toBe('https://api.honcho.dev/v3')
|
||||
})
|
||||
|
||||
test('localhost stays http', () => {
|
||||
expect(normalizeBaseUrl('localhost:8000')).toBe('http://localhost:8000')
|
||||
})
|
||||
})
|
||||
|
||||
describe('resolveConfig', () => {
|
||||
test('host block beats root; env beats host', () => {
|
||||
const file = {
|
||||
workspace: 'root-ws',
|
||||
hosts: { a: { workspace: 'host-ws' } },
|
||||
}
|
||||
expect(resolveConfig(file, { host: 'a', env: emptyEnv }).workspace).toBe('host-ws')
|
||||
expect(
|
||||
resolveConfig(file, { host: 'a', env: { HONCHO_WORKSPACE: 'env-ws' } }).workspace
|
||||
).toBe('env-ws')
|
||||
})
|
||||
|
||||
test('root apiKey / workspaceId aliases still resolve', () => {
|
||||
const cfg = resolveConfig(
|
||||
{ apiKey: 'hch_x', workspaceId: 'from-id' },
|
||||
{ host: 'a', env: emptyEnv }
|
||||
)
|
||||
expect(cfg.apiKey).toBe('hch_x')
|
||||
expect(cfg.workspace).toBe('from-id')
|
||||
})
|
||||
|
||||
test('v1 leftover environmentUrl is ignored', () => {
|
||||
const cfg = resolveConfig(
|
||||
{ schemaVersion: 1, baseUrl: 'https://keep.example', environmentUrl: 'https://old.example' },
|
||||
{ host: 'a', env: emptyEnv }
|
||||
)
|
||||
expect(cfg.baseUrl).toBe('https://keep.example')
|
||||
})
|
||||
|
||||
test('overlay sits below env', () => {
|
||||
expect(
|
||||
resolveConfig(
|
||||
{},
|
||||
{ host: 'a', overlay: { workspace: 'from-overlay' }, env: { HONCHO_WORKSPACE: 'from-env' } }
|
||||
).workspace
|
||||
).toBe('from-env')
|
||||
expect(
|
||||
resolveConfig({}, { host: 'a', overlay: { workspace: 'from-overlay' }, env: emptyEnv }).workspace
|
||||
).toBe('from-overlay')
|
||||
})
|
||||
|
||||
test('empty file uses built-ins; host name is not rewritten', () => {
|
||||
const cfg = resolveConfig({}, { host: 'my-host', env: emptyEnv })
|
||||
expect(cfg.baseUrl).toBe('https://api.honcho.dev')
|
||||
expect(cfg.timeoutMs).toBe(30_000)
|
||||
expect(cfg.enabled).toBe(true)
|
||||
expect(cfg.host).toBe('my-host')
|
||||
expect(cfg.workspace).toBe('my-host')
|
||||
})
|
||||
})
|
||||
|
|
@ -0,0 +1,56 @@
|
|||
import { describe, expect, test } from 'bun:test'
|
||||
import {
|
||||
HEADER_AGENT_MODEL,
|
||||
HEADER_HOST,
|
||||
HEADER_PLUGIN,
|
||||
HEADER_RUNTIME,
|
||||
setTelemetryHeaders,
|
||||
telemetryHeaders,
|
||||
version,
|
||||
} from '../src/index.ts'
|
||||
|
||||
describe('telemetryHeaders', () => {
|
||||
test('empty identity still sends the runtime version', () => {
|
||||
expect(telemetryHeaders()).toEqual({ [HEADER_RUNTIME]: version })
|
||||
})
|
||||
|
||||
test('maps identity to headers', () => {
|
||||
expect(
|
||||
telemetryHeaders({
|
||||
host: 'opencode',
|
||||
hostVersion: '1.3.13',
|
||||
pluginVersion: '0.1.3',
|
||||
model: 'claude-sonnet-4-5',
|
||||
})
|
||||
).toEqual({
|
||||
[HEADER_RUNTIME]: version,
|
||||
[HEADER_HOST]: 'opencode/1.3.13',
|
||||
[HEADER_PLUGIN]: '0.1.3',
|
||||
[HEADER_AGENT_MODEL]: 'claude-sonnet-4-5',
|
||||
})
|
||||
})
|
||||
|
||||
test('merges extra headers last, skipping blanks', () => {
|
||||
const headers = telemetryHeaders({ host: 'codex', pluginVersion: '0.1.1' }, {
|
||||
'X-Custom': 'yes',
|
||||
[HEADER_PLUGIN]: 'override',
|
||||
'X-Empty': ' ',
|
||||
})
|
||||
expect(headers[HEADER_HOST]).toBe('codex')
|
||||
expect(headers[HEADER_PLUGIN]).toBe('override')
|
||||
expect(headers['X-Custom']).toBe('yes')
|
||||
expect(headers).not.toHaveProperty('X-Empty')
|
||||
})
|
||||
})
|
||||
|
||||
describe('setTelemetryHeaders', () => {
|
||||
test('mutates an existing header map in place', () => {
|
||||
const headers = telemetryHeaders({ host: 'cursor', pluginVersion: '0.1.2' })
|
||||
const returned = setTelemetryHeaders(headers, { model: 'claude-opus-4' })
|
||||
expect(returned).toBe(headers)
|
||||
expect(headers[HEADER_HOST]).toBe('cursor')
|
||||
expect(headers[HEADER_PLUGIN]).toBe('0.1.2')
|
||||
expect(headers[HEADER_RUNTIME]).toBe(version)
|
||||
expect(headers[HEADER_AGENT_MODEL]).toBe('claude-opus-4')
|
||||
})
|
||||
})
|
||||
|
|
@ -0,0 +1,13 @@
|
|||
{
|
||||
"compilerOptions": {
|
||||
"target": "ES2022",
|
||||
"module": "ESNext",
|
||||
"moduleResolution": "bundler",
|
||||
"allowImportingTsExtensions": true,
|
||||
"noEmit": true,
|
||||
"strict": true,
|
||||
"skipLibCheck": true,
|
||||
"types": ["node"]
|
||||
},
|
||||
"include": ["src/**/*.ts"]
|
||||
}
|
||||
|
|
@ -0,0 +1,7 @@
|
|||
node_modules
|
||||
.wrangler
|
||||
.dev.vars
|
||||
.env
|
||||
.env.*
|
||||
*.log
|
||||
dist
|
||||
|
|
@ -0,0 +1,19 @@
|
|||
FROM oven/bun:1.2
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
RUN chown bun:bun /app
|
||||
USER bun
|
||||
|
||||
COPY --chown=bun:bun package.json bun.lock bunfig.toml tsconfig.json ./
|
||||
COPY --chown=bun:bun instructions.md ./
|
||||
COPY --chown=bun:bun src ./src
|
||||
|
||||
RUN bun install --frozen-lockfile --production
|
||||
|
||||
EXPOSE 3000
|
||||
|
||||
ENV PORT=3000
|
||||
ENV HOST=0.0.0.0
|
||||
|
||||
CMD ["bun", "src/http.ts"]
|
||||
|
|
@ -1,6 +1,6 @@
|
|||
# Honcho MCP Server
|
||||
|
||||
A Cloudflare Worker that implements the [Model Context Protocol (MCP)](https://modelcontextprotocol.io/) for [Honcho](https://honcho.dev), providing AI memory and personalization tools to LLM clients like Claude Desktop.
|
||||
A [Model Context Protocol (MCP)](https://modelcontextprotocol.io/) server for [Honcho](https://honcho.dev). The hosted path is a Cloudflare Worker; the same tools also run over stdio and over Streamable HTTP (`bun src/http.ts`) for Docker and other long-lived process hosts.
|
||||
|
||||
## Quickstart: Use the Hosted Server
|
||||
|
||||
|
|
@ -45,6 +45,8 @@ Every workspace-scoped tool takes a `workspace_id` argument. If you set `X-Honch
|
|||
```
|
||||
src/
|
||||
index.ts # Worker entry point — parse config, delegate to MCP handler
|
||||
stdio.ts # Local stdio host (bun src/stdio.ts)
|
||||
http.ts # Streamable HTTP host (bun src/http.ts / Docker)
|
||||
server.ts # createServer() — registers all tools on an McpServer
|
||||
config.ts # HonchoConfig, parseConfig(), createClientFactory()
|
||||
types.ts # ToolContext, result helpers
|
||||
|
|
@ -64,25 +66,75 @@ Built on:
|
|||
|
||||
## Self-Hosted Honcho
|
||||
|
||||
If you run Honcho yourself (for privacy, latency, or offline use), deploy the
|
||||
MCP Worker alongside your instance and set `HONCHO_API_URL` in its
|
||||
environment.
|
||||
If you run Honcho yourself, point this server at it with `HONCHO_API_URL`.
|
||||
When unset, requests go to `https://api.honcho.dev`.
|
||||
|
||||
**Local dev (`bun run dev`):** create `mcp/.dev.vars`:
|
||||
**Cloudflare Worker (`bun run dev` / `bun run deploy`):** create `mcp/.dev.vars`:
|
||||
|
||||
```
|
||||
HONCHO_API_URL=http://127.0.0.1:28000
|
||||
```
|
||||
|
||||
**Deployed Worker:**
|
||||
For a deployed Worker: `wrangler secret put HONCHO_API_URL`.
|
||||
|
||||
## HTTP host
|
||||
|
||||
For Docker or any platform that runs a long-lived process, use the Streamable
|
||||
HTTP entry instead of the Worker. Clients keep the same `mcp-remote` shape as
|
||||
`https://mcp.honcho.dev`. Sessions live in process memory — run one instance.
|
||||
|
||||
```bash
|
||||
wrangler secret put HONCHO_API_URL
|
||||
# paste your URL when prompted
|
||||
cd mcp && bun install
|
||||
HONCHO_API_URL=http://127.0.0.1:8000 bun run http
|
||||
```
|
||||
|
||||
When `HONCHO_API_URL` is unset the Worker routes to `https://api.honcho.dev`,
|
||||
so this change is backward-compatible.
|
||||
```bash
|
||||
bunx mcp-remote http://127.0.0.1:3000 \
|
||||
--header "Authorization:Bearer <key>"
|
||||
```
|
||||
|
||||
Auth is the `Authorization: Bearer` header (same as the Worker). Established
|
||||
sessions still require that same bearer. Optional `X-Honcho-Workspace-ID`
|
||||
fills `workspace_id` when the tool argument is omitted.
|
||||
|
||||
`HOST` defaults to `0.0.0.0`, `PORT` to `3000`. `GET /health` is unauthenticated.
|
||||
MCP is served at `/` and `/mcp`. Idle sessions expire after
|
||||
`MCP_SESSION_IDLE_MS` (default 30 minutes); `MCP_SESSION_MAX` (default 128)
|
||||
caps concurrent sessions.
|
||||
|
||||
A platform start command is `bun src/http.ts` (or `bun run http` from `mcp/`).
|
||||
This repo does not ship a `vercel.json`; serverless replicas do not share the
|
||||
in-memory session map.
|
||||
|
||||
### Docker
|
||||
|
||||
```bash
|
||||
docker build -f mcp/Dockerfile -t honcho-mcp mcp
|
||||
docker run --rm -p 3000:3000 \
|
||||
-e HONCHO_API_URL=http://host.docker.internal:8000 \
|
||||
honcho-mcp
|
||||
```
|
||||
|
||||
`docker-compose.yml.example` includes an `mcp` service beside `api` and
|
||||
`deriver` (`HONCHO_API_URL=http://api:8000`, port `127.0.0.1:3000`).
|
||||
|
||||
## Local stdio
|
||||
|
||||
For a local Honcho instance, or any MCP client that spawns a process, run the
|
||||
stdio host. `--cwd` loads `mcp/bunfig.toml` (Markdown loader) from this package.
|
||||
|
||||
```bash
|
||||
cd mcp && bun install
|
||||
|
||||
claude mcp add honcho \
|
||||
-e HONCHO_API_KEY=hch-your-key-here \
|
||||
-e HONCHO_API_URL=http://127.0.0.1:28000 \
|
||||
-e HONCHO_WORKSPACE_ID=my-workspace \
|
||||
-- bun --cwd "$(pwd)" src/stdio.ts
|
||||
```
|
||||
|
||||
`HONCHO_API_URL` defaults to `https://api.honcho.dev`. `HONCHO_WORKSPACE_ID` is
|
||||
optional; without it, pass `workspace_id` on each tool call.
|
||||
|
||||
## Development
|
||||
|
||||
|
|
@ -106,6 +158,8 @@ bun run tsc --noEmit
|
|||
|
||||
### Test locally
|
||||
|
||||
Worker (`bun dev`, port 8787) or HTTP host (`bun run http`, port 3000):
|
||||
|
||||
```bash
|
||||
bunx mcp-remote http://localhost:8787 \
|
||||
--header "Authorization:Bearer <key>"
|
||||
|
|
|
|||
|
|
@ -0,0 +1,5 @@
|
|||
[loader]
|
||||
".md" = "text"
|
||||
|
||||
[run]
|
||||
silent = true
|
||||
|
|
@ -1,7 +1,7 @@
|
|||
{
|
||||
"name": "honcho-mcp",
|
||||
"version": "3.0.0",
|
||||
"description": "Honcho MCP Server — Cloudflare Worker",
|
||||
"description": "Honcho MCP Server",
|
||||
"main": "src/index.ts",
|
||||
"packageManager": "bun@1.2.0",
|
||||
"engines": {
|
||||
|
|
@ -11,6 +11,8 @@
|
|||
"scripts": {
|
||||
"preinstall": "node -e \"const ua=process.env.npm_config_user_agent||'';if(ua.includes('npm')&&!ua.includes('bun')){console.error('❌ Please use bun instead of npm!\\n📦 Run: bun install\\n🌐 Install bun: https://bun.sh/');process.exit(1)}\"",
|
||||
"dev": "wrangler dev",
|
||||
"stdio": "bun src/stdio.ts",
|
||||
"http": "bun src/http.ts",
|
||||
"deploy": "wrangler deploy",
|
||||
"deploy:staging": "wrangler deploy --env staging"
|
||||
},
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ import { Honcho } from "@honcho-ai/sdk";
|
|||
export interface HonchoConfig {
|
||||
apiKey: string;
|
||||
baseUrl: string;
|
||||
/** From X-Honcho-Workspace-ID when set. */
|
||||
/** From X-Honcho-Workspace-ID (HTTP) or HONCHO_WORKSPACE_ID (stdio). */
|
||||
workspaceId?: string;
|
||||
}
|
||||
|
||||
|
|
@ -12,6 +12,12 @@ export interface Env {
|
|||
ALERT_WEBHOOK_URL?: string;
|
||||
}
|
||||
|
||||
export interface EnvConfig {
|
||||
HONCHO_API_KEY?: string;
|
||||
HONCHO_API_URL?: string;
|
||||
HONCHO_WORKSPACE_ID?: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Parse configuration from request headers and Worker env bindings.
|
||||
* Throws only when the Authorization bearer token is missing/empty.
|
||||
|
|
@ -48,8 +54,23 @@ export function parseConfig(request: Request, env: Env = {}): HonchoConfig {
|
|||
};
|
||||
}
|
||||
|
||||
/** Parse configuration from process env. */
|
||||
export function parseEnvConfig(env: EnvConfig): HonchoConfig {
|
||||
const apiKey = env.HONCHO_API_KEY?.trim();
|
||||
if (!apiKey) {
|
||||
throw new Error(
|
||||
"Missing HONCHO_API_KEY. Set HONCHO_API_KEY to your Honcho API key.",
|
||||
);
|
||||
}
|
||||
return {
|
||||
apiKey,
|
||||
baseUrl: env.HONCHO_API_URL?.trim() || "https://api.honcho.dev",
|
||||
workspaceId: env.HONCHO_WORKSPACE_ID?.trim() || undefined,
|
||||
};
|
||||
}
|
||||
|
||||
export const MISSING_WORKSPACE_ID_MESSAGE =
|
||||
"Missing workspace_id. Pass workspace_id on the next tool call, or set the X-Honcho-Workspace-ID header on the connection so it is used automatically.";
|
||||
"Missing workspace_id. Pass workspace_id on the next tool call, or set X-Honcho-Workspace-ID (HTTP) / HONCHO_WORKSPACE_ID (stdio).";
|
||||
|
||||
export function resolveWorkspaceId(
|
||||
config: HonchoConfig,
|
||||
|
|
|
|||
|
|
@ -0,0 +1,56 @@
|
|||
import { expect, test } from "bun:test";
|
||||
import { fetch } from "./http.ts";
|
||||
|
||||
const origin = "http://127.0.0.1:3000";
|
||||
|
||||
const initializeBody = {
|
||||
jsonrpc: "2.0",
|
||||
id: 1,
|
||||
method: "initialize",
|
||||
params: {
|
||||
protocolVersion: "2024-11-05",
|
||||
capabilities: {},
|
||||
clientInfo: { name: "test", version: "0.0.0" },
|
||||
},
|
||||
};
|
||||
|
||||
const pingBody = { jsonrpc: "2.0", id: 2, method: "ping" };
|
||||
|
||||
function mcpPost(headers: Record<string, string>, body: unknown) {
|
||||
return fetch(
|
||||
new Request(`${origin}/mcp`, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
Accept: "application/json, text/event-stream",
|
||||
"Content-Type": "application/json",
|
||||
...headers,
|
||||
},
|
||||
body: JSON.stringify(body),
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
test("established sessions require the initialize bearer", async () => {
|
||||
const init = await mcpPost(
|
||||
{ Authorization: "Bearer key-a" },
|
||||
initializeBody,
|
||||
);
|
||||
expect(init.status).toBe(200);
|
||||
const sessionId = init.headers.get("mcp-session-id");
|
||||
expect(sessionId).toBeTruthy();
|
||||
|
||||
const missing = await mcpPost({ "mcp-session-id": sessionId! }, pingBody);
|
||||
expect(missing.status).toBe(401);
|
||||
|
||||
const wrong = await mcpPost(
|
||||
{ Authorization: "Bearer key-b", "mcp-session-id": sessionId! },
|
||||
pingBody,
|
||||
);
|
||||
expect(wrong.status).toBe(401);
|
||||
|
||||
const ok = await mcpPost(
|
||||
{ Authorization: "Bearer key-a", "mcp-session-id": sessionId! },
|
||||
pingBody,
|
||||
);
|
||||
expect(ok.status).toBe(200);
|
||||
});
|
||||
|
|
@ -0,0 +1,281 @@
|
|||
import { WebStandardStreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/webStandardStreamableHttp.js";
|
||||
import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js";
|
||||
import {
|
||||
createClientFactory,
|
||||
createUnscopedClient,
|
||||
parseConfig,
|
||||
type Env,
|
||||
type HonchoConfig,
|
||||
} from "./config.js";
|
||||
import { createServer } from "./server.js";
|
||||
import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
|
||||
|
||||
declare const process: {
|
||||
env: Record<string, string | undefined>;
|
||||
};
|
||||
|
||||
declare const Bun: {
|
||||
serve(options: {
|
||||
hostname: string;
|
||||
port: number;
|
||||
fetch(request: Request): Response | Promise<Response>;
|
||||
}): { hostname: string; port: number };
|
||||
};
|
||||
|
||||
const CORS_ORIGIN = "*";
|
||||
const CORS_METHODS = "GET, POST, DELETE, OPTIONS";
|
||||
const CORS_ALLOWED_HEADERS =
|
||||
"Content-Type, Authorization, X-Honcho-Workspace-ID, mcp-session-id, mcp-protocol-version, last-event-id";
|
||||
|
||||
const CORS_HEADERS: Record<string, string> = {
|
||||
"Access-Control-Allow-Origin": CORS_ORIGIN,
|
||||
"Access-Control-Allow-Methods": CORS_METHODS,
|
||||
"Access-Control-Allow-Headers": CORS_ALLOWED_HEADERS,
|
||||
"Access-Control-Expose-Headers": "WWW-Authenticate, mcp-session-id",
|
||||
};
|
||||
|
||||
const PROTECTED_RESOURCE_PATH = "/.well-known/oauth-protected-resource";
|
||||
const MCP_PATHS = new Set(["/", "/mcp"]);
|
||||
|
||||
type Session = {
|
||||
transport: WebStandardStreamableHTTPServerTransport;
|
||||
server: McpServer;
|
||||
lastSeen: number;
|
||||
apiKey: string;
|
||||
};
|
||||
|
||||
const sessions = new Map<string, Session>();
|
||||
const DEFAULT_SESSION_IDLE_MS = 30 * 60 * 1000;
|
||||
const DEFAULT_SESSION_MAX = 128;
|
||||
|
||||
function envInt(name: string, fallback: number): number {
|
||||
const n = Number(process.env[name]);
|
||||
return Number.isFinite(n) && n > 0 ? n : fallback;
|
||||
}
|
||||
|
||||
function dropSession(id: string): void {
|
||||
const session = sessions.get(id);
|
||||
if (!session) return;
|
||||
sessions.delete(id);
|
||||
void session.transport.close();
|
||||
void session.server.close();
|
||||
}
|
||||
|
||||
function sweepSessions(): void {
|
||||
const idleMs = envInt("MCP_SESSION_IDLE_MS", DEFAULT_SESSION_IDLE_MS);
|
||||
const now = Date.now();
|
||||
for (const [id, session] of sessions) {
|
||||
if (now - session.lastSeen > idleMs) dropSession(id);
|
||||
}
|
||||
}
|
||||
|
||||
function envBindings(): Env {
|
||||
return { HONCHO_API_URL: process.env.HONCHO_API_URL };
|
||||
}
|
||||
|
||||
function authorizationServer(): string {
|
||||
return process.env.HONCHO_API_URL?.trim() || "https://api.honcho.dev";
|
||||
}
|
||||
|
||||
function withCors(response: Response): Response {
|
||||
const headers = new Headers(response.headers);
|
||||
for (const [key, value] of Object.entries(CORS_HEADERS)) {
|
||||
headers.set(key, value);
|
||||
}
|
||||
return new Response(response.body, {
|
||||
status: response.status,
|
||||
statusText: response.statusText,
|
||||
headers,
|
||||
});
|
||||
}
|
||||
|
||||
function jsonResponse(
|
||||
body: unknown,
|
||||
status: number,
|
||||
extraHeaders?: Record<string, string>,
|
||||
): Response {
|
||||
return new Response(JSON.stringify(body), {
|
||||
status,
|
||||
headers: {
|
||||
"Content-Type": "application/json",
|
||||
...CORS_HEADERS,
|
||||
...extraHeaders,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
function configForRequest(request: Request) {
|
||||
return parseConfig(request, envBindings());
|
||||
}
|
||||
|
||||
function configOrUnauthorized(request: Request): HonchoConfig | Response {
|
||||
try {
|
||||
return configForRequest(request);
|
||||
} catch (e) {
|
||||
const message = e instanceof Error ? e.message : "Invalid request";
|
||||
return unauthorized(request, message);
|
||||
}
|
||||
}
|
||||
|
||||
function unauthorized(request: Request, message: string): Response {
|
||||
const resourceMetadata = `${new URL(request.url).origin}${PROTECTED_RESOURCE_PATH}`;
|
||||
return jsonResponse(
|
||||
{ error: message },
|
||||
401,
|
||||
{
|
||||
"WWW-Authenticate": `Bearer resource_metadata="${resourceMetadata}"`,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
async function handleMcp(request: Request): Promise<Response> {
|
||||
sweepSessions();
|
||||
const sessionId = request.headers.get("mcp-session-id");
|
||||
if (sessionId) {
|
||||
const existing = sessions.get(sessionId);
|
||||
if (existing) {
|
||||
const config = configOrUnauthorized(request);
|
||||
if (config instanceof Response) return config;
|
||||
if (config.apiKey !== existing.apiKey) {
|
||||
return unauthorized(
|
||||
request,
|
||||
"Authorization does not match this session.",
|
||||
);
|
||||
}
|
||||
existing.lastSeen = Date.now();
|
||||
return withCors(await existing.transport.handleRequest(request));
|
||||
}
|
||||
}
|
||||
|
||||
if (request.method !== "POST") {
|
||||
return jsonResponse(
|
||||
{
|
||||
jsonrpc: "2.0",
|
||||
error: {
|
||||
code: -32000,
|
||||
message: "Bad Request: No valid session ID provided",
|
||||
},
|
||||
id: null,
|
||||
},
|
||||
400,
|
||||
);
|
||||
}
|
||||
|
||||
let body: unknown;
|
||||
try {
|
||||
body = await request.json();
|
||||
} catch {
|
||||
return jsonResponse(
|
||||
{
|
||||
jsonrpc: "2.0",
|
||||
error: { code: -32700, message: "Parse error: Invalid JSON" },
|
||||
id: null,
|
||||
},
|
||||
400,
|
||||
);
|
||||
}
|
||||
|
||||
const messages = Array.isArray(body) ? body : [body];
|
||||
if (!messages.some((message) => isInitializeRequest(message))) {
|
||||
return jsonResponse(
|
||||
{
|
||||
jsonrpc: "2.0",
|
||||
error: {
|
||||
code: -32000,
|
||||
message: "Bad Request: No valid session ID provided",
|
||||
},
|
||||
id: null,
|
||||
},
|
||||
400,
|
||||
);
|
||||
}
|
||||
|
||||
const config = configOrUnauthorized(request);
|
||||
if (config instanceof Response) return config;
|
||||
|
||||
const server = createServer({
|
||||
config,
|
||||
clientFor: createClientFactory(config),
|
||||
unscoped: createUnscopedClient(config),
|
||||
});
|
||||
|
||||
const maxSessions = envInt("MCP_SESSION_MAX", DEFAULT_SESSION_MAX);
|
||||
if (sessions.size >= maxSessions) {
|
||||
return jsonResponse(
|
||||
{
|
||||
jsonrpc: "2.0",
|
||||
error: {
|
||||
code: -32000,
|
||||
message: "Too many active sessions",
|
||||
},
|
||||
id: null,
|
||||
},
|
||||
503,
|
||||
);
|
||||
}
|
||||
|
||||
const transport = new WebStandardStreamableHTTPServerTransport({
|
||||
sessionIdGenerator: () => crypto.randomUUID(),
|
||||
onsessioninitialized: (id) => {
|
||||
sessions.set(id, {
|
||||
transport,
|
||||
server,
|
||||
lastSeen: Date.now(),
|
||||
apiKey: config.apiKey,
|
||||
});
|
||||
},
|
||||
});
|
||||
transport.onclose = () => {
|
||||
const id = transport.sessionId;
|
||||
if (id) sessions.delete(id);
|
||||
};
|
||||
|
||||
await server.connect(transport);
|
||||
return withCors(
|
||||
await transport.handleRequest(request, { parsedBody: body }),
|
||||
);
|
||||
}
|
||||
|
||||
export async function fetch(request: Request): Promise<Response> {
|
||||
if (request.method === "OPTIONS") {
|
||||
return new Response(null, { status: 204, headers: CORS_HEADERS });
|
||||
}
|
||||
|
||||
const pathname = new URL(request.url).pathname;
|
||||
|
||||
if (pathname === "/health") {
|
||||
return jsonResponse({ status: "ok" }, 200);
|
||||
}
|
||||
|
||||
if (pathname === PROTECTED_RESOURCE_PATH) {
|
||||
return jsonResponse(
|
||||
{
|
||||
resource: new URL(request.url).origin,
|
||||
authorization_servers: [authorizationServer()],
|
||||
bearer_methods_supported: ["header"],
|
||||
scopes_supported: ["read", "write"],
|
||||
},
|
||||
200,
|
||||
);
|
||||
}
|
||||
|
||||
if (!MCP_PATHS.has(pathname)) {
|
||||
return jsonResponse({ error: "Not Found" }, 404);
|
||||
}
|
||||
|
||||
try {
|
||||
return await handleMcp(request);
|
||||
} catch (e) {
|
||||
const message =
|
||||
e instanceof Error ? e.message : "Internal server error";
|
||||
return jsonResponse({ error: message }, 500);
|
||||
}
|
||||
}
|
||||
|
||||
const isMain = Boolean((import.meta as { main?: boolean }).main);
|
||||
if (isMain) {
|
||||
const hostname = process.env.HOST?.trim() || "0.0.0.0";
|
||||
const port = Number(process.env.PORT) || 3000;
|
||||
Bun.serve({ hostname, port, fetch });
|
||||
console.error(`honcho-mcp listening on http://${hostname}:${port}`);
|
||||
}
|
||||
|
|
@ -0,0 +1,30 @@
|
|||
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
|
||||
import {
|
||||
createClientFactory,
|
||||
createUnscopedClient,
|
||||
parseEnvConfig,
|
||||
} from "./config.js";
|
||||
import { createServer } from "./server.js";
|
||||
|
||||
declare const process: {
|
||||
env: Record<string, string | undefined>;
|
||||
exit(code?: number): never;
|
||||
};
|
||||
|
||||
try {
|
||||
const config = parseEnvConfig({
|
||||
HONCHO_API_KEY: process.env.HONCHO_API_KEY,
|
||||
HONCHO_API_URL: process.env.HONCHO_API_URL,
|
||||
HONCHO_WORKSPACE_ID: process.env.HONCHO_WORKSPACE_ID,
|
||||
});
|
||||
const server = createServer({
|
||||
config,
|
||||
clientFor: createClientFactory(config),
|
||||
unscoped: createUnscopedClient(config),
|
||||
});
|
||||
await server.connect(new StdioServerTransport());
|
||||
} catch (e) {
|
||||
const message = e instanceof Error ? e.message : String(e);
|
||||
console.error(message);
|
||||
process.exit(1);
|
||||
}
|
||||
|
|
@ -10,5 +10,5 @@
|
|||
"types": ["@cloudflare/workers-types"]
|
||||
},
|
||||
"include": ["src/**/*.ts"],
|
||||
"exclude": ["node_modules"]
|
||||
"exclude": ["node_modules", "src/**/*.test.ts"]
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "honcho"
|
||||
version = "3.1.0"
|
||||
version = "3.1.1"
|
||||
description = "Honcho Server"
|
||||
authors = [
|
||||
{name = "Plastic Labs", email = "hello@plasticlabs.ai"},
|
||||
|
|
|
|||
|
|
@ -0,0 +1,43 @@
|
|||
{
|
||||
"$schema": "https://json-schema.org/draft/2020-12/schema",
|
||||
"$id": "https://spec.honcho.dev/config/v1.json",
|
||||
"type": "object",
|
||||
"additionalProperties": true,
|
||||
"$defs": {
|
||||
"oauth": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"accessToken": { "type": "string" },
|
||||
"refreshToken": { "type": "string" },
|
||||
"expiresAt": { "type": "string" }
|
||||
}
|
||||
},
|
||||
"auth": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"apiKey": { "type": "string" },
|
||||
"oauth": { "$ref": "#/$defs/oauth" }
|
||||
}
|
||||
},
|
||||
"hostBlock": {
|
||||
"type": "object",
|
||||
"additionalProperties": true,
|
||||
"properties": {
|
||||
"peerName": { "type": "string" },
|
||||
"workspace": { "type": "string" },
|
||||
"baseUrl": { "type": "string" },
|
||||
"timeoutMs": { "type": "number" },
|
||||
"enabled": { "type": "boolean" },
|
||||
"auth": { "$ref": "#/$defs/auth" }
|
||||
}
|
||||
}
|
||||
},
|
||||
"allOf": [{ "$ref": "#/$defs/hostBlock" }],
|
||||
"properties": {
|
||||
"schemaVersion": { "type": "integer", "const": 1 },
|
||||
"hosts": {
|
||||
"type": "object",
|
||||
"additionalProperties": { "$ref": "#/$defs/hostBlock" }
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file.
|
|||
The format is based on [Keep a Changelog](http://keepachangelog.com/)
|
||||
and this project adheres to [Semantic Versioning](http://semver.org/).
|
||||
|
||||
## [Unreleased]
|
||||
|
||||
### Added
|
||||
|
||||
- Optional per-call `timeout` on synchronous and asynchronous `Peer.chat()`. It overrides the timeout for each HTTP attempt; when omitted or set to `None`, the client-wide timeout configured on `Honcho` remains in effect.
|
||||
|
||||
## [2.4.0] - 2026-08-25
|
||||
|
||||
### Added
|
||||
|
|
|
|||
|
|
@ -777,6 +777,7 @@ class PeerAio(AsyncMetadataConfigMixin):
|
|||
reasoning_level: Literal["minimal", "low", "medium", "high", "max"]
|
||||
| None = None,
|
||||
response_format: type[TResponseFormat],
|
||||
timeout: float | None = None,
|
||||
) -> TResponseFormat | None: ...
|
||||
|
||||
@overload
|
||||
|
|
@ -791,6 +792,7 @@ class PeerAio(AsyncMetadataConfigMixin):
|
|||
reasoning_level: Literal["minimal", "low", "medium", "high", "max"]
|
||||
| None = None,
|
||||
response_format: dict[str, Any] | None = None,
|
||||
timeout: float | None = None,
|
||||
) -> str | None: ...
|
||||
|
||||
@validate_call(config=ConfigDict(arbitrary_types_allowed=True))
|
||||
|
|
@ -805,12 +807,17 @@ class PeerAio(AsyncMetadataConfigMixin):
|
|||
reasoning_level: Literal["minimal", "low", "medium", "high", "max"]
|
||||
| None = None,
|
||||
response_format: type[BaseModel] | dict[str, Any] | None = None,
|
||||
timeout: float | None = Field(
|
||||
None, gt=0, description="Timeout in seconds for this chat request"
|
||||
),
|
||||
) -> BaseModel | str | None:
|
||||
"""Query the peer's representation asynchronously.
|
||||
|
||||
See Peer.chat for parameter details. When response_format is a Pydantic
|
||||
model class, the answer is parsed into an instance of it; when it is a
|
||||
JSON Schema dict, the answer is a JSON string.
|
||||
JSON Schema dict, the answer is a JSON string. When timeout is omitted,
|
||||
the Honcho client's configured timeout is used; retries can extend total
|
||||
elapsed time.
|
||||
"""
|
||||
await self._peer._honcho._ensure_workspace_async()
|
||||
target_id = resolve_id(target)
|
||||
|
|
@ -835,6 +842,7 @@ class PeerAio(AsyncMetadataConfigMixin):
|
|||
data = await self._peer._honcho._async_http_client.post(
|
||||
routes.peer_chat(self._peer.workspace_id, self._peer.id),
|
||||
body=body,
|
||||
timeout=timeout,
|
||||
)
|
||||
content = data.get("content")
|
||||
if not content:
|
||||
|
|
|
|||
|
|
@ -246,6 +246,7 @@ class Peer(PeerBase, MetadataConfigMixin):
|
|||
reasoning_level: Literal["minimal", "low", "medium", "high", "max"]
|
||||
| None = None,
|
||||
response_format: type[TResponseFormat],
|
||||
timeout: float | None = None,
|
||||
) -> TResponseFormat | None: ...
|
||||
|
||||
@overload
|
||||
|
|
@ -260,6 +261,7 @@ class Peer(PeerBase, MetadataConfigMixin):
|
|||
reasoning_level: Literal["minimal", "low", "medium", "high", "max"]
|
||||
| None = None,
|
||||
response_format: dict[str, Any] | None = None,
|
||||
timeout: float | None = None,
|
||||
) -> str | None: ...
|
||||
|
||||
@validate_call(config=ConfigDict(arbitrary_types_allowed=True))
|
||||
|
|
@ -274,6 +276,9 @@ class Peer(PeerBase, MetadataConfigMixin):
|
|||
reasoning_level: Literal["minimal", "low", "medium", "high", "max"]
|
||||
| None = None,
|
||||
response_format: type[BaseModel] | dict[str, Any] | None = None,
|
||||
timeout: float | None = Field(
|
||||
None, gt=0, description="Timeout in seconds for this chat request"
|
||||
),
|
||||
) -> BaseModel | str | None:
|
||||
"""
|
||||
Query the peer's representation with a natural language question.
|
||||
|
|
@ -310,6 +315,9 @@ class Peer(PeerBase, MetadataConfigMixin):
|
|||
model class to get a parsed instance back, or a raw
|
||||
JSON Schema dict (root type "object") to get the
|
||||
answer as a JSON string.
|
||||
timeout: Optional timeout in seconds for each HTTP attempt made by
|
||||
this request. When omitted, the Honcho client's configured
|
||||
timeout is used. Retries can extend total elapsed time.
|
||||
|
||||
Returns:
|
||||
Response string containing the answer (a JSON string when a schema
|
||||
|
|
@ -342,6 +350,7 @@ class Peer(PeerBase, MetadataConfigMixin):
|
|||
data = self._honcho._http.post(
|
||||
routes.peer_chat(self.workspace_id, self.id),
|
||||
body=body,
|
||||
timeout=timeout,
|
||||
)
|
||||
content = data.get("content")
|
||||
if not content:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,146 @@
|
|||
"""Read-only polling of the deriver's outstanding work. Schedules nothing."""
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from logging import getLogger
|
||||
|
||||
import sentry_sdk
|
||||
|
||||
from src import crud, schemas
|
||||
from src.config import settings
|
||||
from src.dependencies import tracked_db
|
||||
from src.dreamer.dream_due import count_due_dreams
|
||||
from src.telemetry import prometheus_metrics
|
||||
|
||||
logger = getLogger(__name__)
|
||||
|
||||
|
||||
def active_work_seconds() -> float:
|
||||
"""The value reported when work is ready for a deriver now."""
|
||||
return float(max(settings.DERIVER.REPRESENTATION_BATCH_MAX_AGE_SECONDS, 1))
|
||||
|
||||
|
||||
@dataclass
|
||||
class DeriverMetricsSnapshot:
|
||||
"""The last good poll result, served to callers of the route."""
|
||||
|
||||
signal_seconds: float = 0.0
|
||||
dreams_due: int = 0
|
||||
stats: schemas.DeriverMetrics = field(default_factory=schemas.DeriverMetrics)
|
||||
measured_at: float | None = None
|
||||
|
||||
@property
|
||||
def age_seconds(self) -> float | None:
|
||||
if self.measured_at is None:
|
||||
return None
|
||||
return max(0.0, time.time() - self.measured_at)
|
||||
|
||||
|
||||
def outstanding_work_seconds(
|
||||
stats: schemas.DeriverMetrics, *, dreams_due: int
|
||||
) -> float:
|
||||
"""Seconds of outstanding deriver work, 0 when there is nothing to do."""
|
||||
if (
|
||||
stats.eligible_work_units > 0
|
||||
or stats.claimed_work_units > 0
|
||||
or stats.embeddings_pending_due > 0
|
||||
or dreams_due > 0
|
||||
):
|
||||
return active_work_seconds()
|
||||
if stats.pending_items > 0:
|
||||
return stats.oldest_pending_age_seconds
|
||||
return 0.0
|
||||
|
||||
|
||||
class DeriverMetricsPoller:
|
||||
"""Refreshes the deriver gauges and the cached snapshot on a timer."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._task: asyncio.Task[None] | None = None
|
||||
self._shutdown_event: asyncio.Event = asyncio.Event()
|
||||
self._snapshot: DeriverMetricsSnapshot = DeriverMetricsSnapshot()
|
||||
self._next_dream_poll: float | None = None
|
||||
self._dreams_due: int = 0
|
||||
|
||||
@property
|
||||
def snapshot(self) -> DeriverMetricsSnapshot:
|
||||
return self._snapshot
|
||||
|
||||
async def start(self) -> None:
|
||||
if self._task is not None:
|
||||
logger.warning("DeriverMetricsPoller already running")
|
||||
return
|
||||
self._shutdown_event.clear()
|
||||
self._task = asyncio.create_task(self._loop())
|
||||
logger.info(
|
||||
"DeriverMetricsPoller started, interval %ss",
|
||||
settings.DERIVER.BACKLOG_METRICS_POLL_INTERVAL_SECONDS,
|
||||
)
|
||||
|
||||
async def shutdown(self) -> None:
|
||||
if self._task is None:
|
||||
return
|
||||
logger.info("Shutting down DeriverMetricsPoller...")
|
||||
self._shutdown_event.set()
|
||||
try:
|
||||
await asyncio.wait_for(self._task, timeout=5.0)
|
||||
except TimeoutError:
|
||||
logger.warning("DeriverMetricsPoller shutdown timed out, cancelling task")
|
||||
self._task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await self._task
|
||||
self._task = None
|
||||
logger.info("DeriverMetricsPoller stopped")
|
||||
|
||||
async def _loop(self) -> None:
|
||||
interval = settings.DERIVER.BACKLOG_METRICS_POLL_INTERVAL_SECONDS
|
||||
while not self._shutdown_event.is_set():
|
||||
try:
|
||||
await self.refresh()
|
||||
except Exception as e:
|
||||
logger.error("DeriverMetricsPoller refresh failed: %s", e)
|
||||
if settings.SENTRY.ENABLED:
|
||||
sentry_sdk.capture_exception(e)
|
||||
with contextlib.suppress(TimeoutError):
|
||||
await asyncio.wait_for(self._shutdown_event.wait(), timeout=interval)
|
||||
|
||||
async def refresh(self) -> None:
|
||||
"""One read-only pass. The snapshot only advances on a complete pass."""
|
||||
async with tracked_db("deriver_metrics", read_only=True) as db:
|
||||
stats = await crud.get_deriver_metrics(db)
|
||||
if self._dream_poll_due():
|
||||
self._dreams_due = await count_due_dreams(db)
|
||||
self._next_dream_poll = (
|
||||
time.monotonic() + settings.DREAM.DUE_POLL_INTERVAL_SECONDS
|
||||
)
|
||||
|
||||
signal = outstanding_work_seconds(stats, dreams_due=self._dreams_due)
|
||||
measured_at = time.time()
|
||||
|
||||
self._snapshot = DeriverMetricsSnapshot(
|
||||
signal_seconds=signal,
|
||||
dreams_due=self._dreams_due,
|
||||
stats=stats,
|
||||
measured_at=measured_at,
|
||||
)
|
||||
|
||||
metrics = prometheus_metrics
|
||||
metrics.set_deriver_metrics(
|
||||
eligible_work_units=stats.eligible_work_units,
|
||||
claimed_work_units=stats.claimed_work_units,
|
||||
pending_items=stats.pending_items,
|
||||
oldest_pending_age_seconds=stats.oldest_pending_age_seconds,
|
||||
embeddings_pending=stats.embeddings_pending,
|
||||
embeddings_pending_due=stats.embeddings_pending_due,
|
||||
)
|
||||
metrics.set_dreams_due(count=self._dreams_due)
|
||||
metrics.set_deriver_outstanding_work(seconds=signal)
|
||||
metrics.set_deriver_metrics_last_success(timestamp=measured_at)
|
||||
|
||||
def _dream_poll_due(self) -> bool:
|
||||
"""The dream query is far more expensive, so it runs on its own spacing."""
|
||||
return (
|
||||
self._next_dream_poll is None or time.monotonic() >= self._next_dream_poll
|
||||
)
|
||||
|
|
@ -972,6 +972,8 @@ class DeriverSettings(HonchoSettings):
|
|||
# When enabled, bypasses the batch token threshold and processes work immediately
|
||||
FLUSH_ENABLED: bool = False
|
||||
|
||||
BACKLOG_METRICS_POLL_INTERVAL_SECONDS: Annotated[int, Field(default=30, ge=1)] = 30
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def _merge_model_config_defaults(cls, data: Any) -> Any:
|
||||
|
|
@ -1351,6 +1353,7 @@ class DreamSettings(HonchoSettings):
|
|||
DOCUMENT_THRESHOLD: Annotated[int, Field(default=50, gt=0, le=1000)] = 50
|
||||
IDLE_TIMEOUT_MINUTES: Annotated[int, Field(default=60, gt=0, le=1440)] = 60
|
||||
MIN_HOURS_BETWEEN_DREAMS: Annotated[int, Field(default=8, gt=0, le=72)] = 8
|
||||
DUE_POLL_INTERVAL_SECONDS: Annotated[int, Field(default=300, ge=1)] = 300
|
||||
ENABLED_TYPES: list[str] = ["omni"]
|
||||
|
||||
# Agent iteration limit - increased for extended reasoning workflow
|
||||
|
|
|
|||
|
|
@ -3,7 +3,11 @@ from .collection import (
|
|||
get_or_create_collection,
|
||||
update_collection_internal_metadata,
|
||||
)
|
||||
from .deriver import get_deriver_status, get_queue_status
|
||||
from .deriver import (
|
||||
get_deriver_metrics,
|
||||
get_deriver_status,
|
||||
get_queue_status,
|
||||
)
|
||||
from .document import (
|
||||
CreateDocumentsResult,
|
||||
create_documents,
|
||||
|
|
@ -105,6 +109,7 @@ __all__ = [
|
|||
"get_or_create_collection",
|
||||
"update_collection_internal_metadata",
|
||||
# Deriver
|
||||
"get_deriver_metrics",
|
||||
"get_deriver_status",
|
||||
"get_queue_status",
|
||||
# Document
|
||||
|
|
|
|||
|
|
@ -1,15 +1,165 @@
|
|||
from collections.abc import Sequence
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from logging import getLogger
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import Select, case, func, or_, select
|
||||
from sqlalchemy import ColumnElement, Select, case, func, or_, select
|
||||
from sqlalchemy.engine import Row
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from src import models, schemas
|
||||
from src.config import settings
|
||||
|
||||
logger = getLogger(__name__)
|
||||
|
||||
REPRESENTATION_WORK_UNIT_PREFIX = "representation:"
|
||||
|
||||
|
||||
def representation_batch_threshold_clause(
|
||||
*,
|
||||
work_unit_key: ColumnElement[str],
|
||||
total_tokens: ColumnElement[Any],
|
||||
oldest_created_at: ColumnElement[Any],
|
||||
) -> ColumnElement[bool] | None:
|
||||
"""The batch gate a representation work unit passes before it is claimable, or None when no gate applies."""
|
||||
if settings.DERIVER.FLUSH_ENABLED:
|
||||
return None
|
||||
|
||||
target_tokens = settings.DERIVER.REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS
|
||||
if target_tokens <= 0:
|
||||
return None
|
||||
|
||||
threshold: ColumnElement[bool] = func.coalesce(total_tokens, 0) >= target_tokens
|
||||
|
||||
max_age_seconds = settings.DERIVER.REPRESENTATION_BATCH_MAX_AGE_SECONDS
|
||||
if max_age_seconds > 0:
|
||||
threshold = or_(
|
||||
threshold,
|
||||
oldest_created_at <= func.now() - timedelta(seconds=max_age_seconds),
|
||||
)
|
||||
|
||||
return or_(
|
||||
~work_unit_key.startswith(REPRESENTATION_WORK_UNIT_PREFIX),
|
||||
threshold,
|
||||
)
|
||||
|
||||
|
||||
def unclaimed_work_unit_clause(
|
||||
work_unit_key: ColumnElement[str],
|
||||
) -> ColumnElement[bool]:
|
||||
"""No claim row exists for this work unit, stale ones included."""
|
||||
return (
|
||||
~select(models.ActiveQueueSession.id)
|
||||
.where(models.ActiveQueueSession.work_unit_key == work_unit_key)
|
||||
.exists()
|
||||
)
|
||||
|
||||
|
||||
def stale_claim_cutoff() -> datetime:
|
||||
return datetime.now(UTC) - timedelta(
|
||||
minutes=settings.DERIVER.STALE_SESSION_TIMEOUT_MINUTES
|
||||
)
|
||||
|
||||
|
||||
def not_live_claimed_work_unit_clause(
|
||||
work_unit_key: ColumnElement[str],
|
||||
) -> ColumnElement[bool]:
|
||||
"""No claim refreshed inside the stale timeout exists, so a stale claim leaves its work unit claimable."""
|
||||
return (
|
||||
~select(models.ActiveQueueSession.id)
|
||||
.where(
|
||||
models.ActiveQueueSession.work_unit_key == work_unit_key,
|
||||
models.ActiveQueueSession.last_updated >= stale_claim_cutoff(),
|
||||
)
|
||||
.exists()
|
||||
)
|
||||
|
||||
|
||||
async def get_deriver_metrics(db: AsyncSession) -> schemas.DeriverMetrics:
|
||||
"""Count the outstanding deriver work in the whole database, read-only."""
|
||||
from src.reconciler.sync_vectors import backoff_eligible # noqa: PLC0415
|
||||
|
||||
token_stats = (
|
||||
select(
|
||||
models.QueueItem.work_unit_key,
|
||||
func.sum(models.Message.token_count).label("total_tokens"),
|
||||
func.min(models.QueueItem.created_at).label("oldest_created_at"),
|
||||
)
|
||||
.join(models.Message, models.QueueItem.message_id == models.Message.id)
|
||||
.where(~models.QueueItem.processed)
|
||||
.where(
|
||||
models.QueueItem.work_unit_key.startswith(REPRESENTATION_WORK_UNIT_PREFIX)
|
||||
)
|
||||
.group_by(models.QueueItem.work_unit_key)
|
||||
.subquery()
|
||||
)
|
||||
|
||||
work_units = (
|
||||
select(models.QueueItem.work_unit_key)
|
||||
.where(~models.QueueItem.processed)
|
||||
.group_by(models.QueueItem.work_unit_key)
|
||||
.subquery()
|
||||
)
|
||||
|
||||
eligible = (
|
||||
select(func.count())
|
||||
.select_from(work_units)
|
||||
.outerjoin(
|
||||
token_stats,
|
||||
work_units.c.work_unit_key == token_stats.c.work_unit_key,
|
||||
)
|
||||
.where(not_live_claimed_work_unit_clause(work_units.c.work_unit_key))
|
||||
)
|
||||
|
||||
threshold_clause = representation_batch_threshold_clause(
|
||||
work_unit_key=work_units.c.work_unit_key,
|
||||
total_tokens=token_stats.c.total_tokens,
|
||||
oldest_created_at=token_stats.c.oldest_created_at,
|
||||
)
|
||||
if threshold_clause is not None:
|
||||
eligible = eligible.where(threshold_clause)
|
||||
|
||||
claimed = (
|
||||
select(func.count())
|
||||
.select_from(models.ActiveQueueSession)
|
||||
.where(models.ActiveQueueSession.last_updated >= stale_claim_cutoff())
|
||||
)
|
||||
|
||||
pending = select(
|
||||
func.count(models.QueueItem.id),
|
||||
func.coalesce(
|
||||
func.extract("epoch", func.now() - func.min(models.QueueItem.created_at)),
|
||||
0,
|
||||
),
|
||||
).where(~models.QueueItem.processed)
|
||||
|
||||
embeddings = select(
|
||||
func.count(),
|
||||
func.coalesce(
|
||||
func.sum(
|
||||
case(
|
||||
(backoff_eligible(models.MessageEmbedding.last_sync_at), 1),
|
||||
else_=0,
|
||||
)
|
||||
),
|
||||
0,
|
||||
),
|
||||
).where(models.MessageEmbedding.sync_state == "pending")
|
||||
|
||||
eligible_count = (await db.execute(eligible)).scalar_one()
|
||||
claimed_count = (await db.execute(claimed)).scalar_one()
|
||||
pending_count, oldest_age = (await db.execute(pending)).one()
|
||||
embeddings_pending, embeddings_due = (await db.execute(embeddings)).one()
|
||||
|
||||
return schemas.DeriverMetrics(
|
||||
eligible_work_units=int(eligible_count),
|
||||
claimed_work_units=int(claimed_count),
|
||||
pending_items=int(pending_count),
|
||||
oldest_pending_age_seconds=float(oldest_age),
|
||||
embeddings_pending=int(embeddings_pending),
|
||||
embeddings_pending_due=int(embeddings_due),
|
||||
)
|
||||
|
||||
|
||||
async def get_queue_status(
|
||||
db: AsyncSession,
|
||||
|
|
|
|||
|
|
@ -1,13 +1,14 @@
|
|||
import asyncio
|
||||
import datetime
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from logging import getLogger
|
||||
from typing import Any, cast
|
||||
from typing import Any, Literal, cast
|
||||
|
||||
from sqlalchemy import delete, select, update
|
||||
from sqlalchemy.engine import CursorResult
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
from sqlalchemy.exc import DBAPIError, IntegrityError, SQLAlchemyError
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.sql import Select
|
||||
from sqlalchemy.sql.functions import func
|
||||
|
|
@ -210,6 +211,24 @@ def _uses_pgvector() -> bool:
|
|||
)
|
||||
|
||||
|
||||
# Shared by is_rejected_duplicate and create_documents candidate resolution.
|
||||
_SEMANTIC_DUP_MAX_DISTANCE = 0.05
|
||||
_SEMANTIC_DUP_TOP_K = 1
|
||||
_SEMANTIC_CANDIDATE_CONCURRENCY = 8
|
||||
|
||||
|
||||
def _semantic_dup_filters(doc: schemas.DocumentCreate) -> dict[str, Any] | None:
|
||||
"""Merge scope for semantic dedup: never across levels, never across
|
||||
sessions for explicit documents. None when the document has no valid
|
||||
merge partner (session-less explicit)."""
|
||||
filters: dict[str, Any] = {"level": doc.level}
|
||||
if doc.level == "explicit":
|
||||
if doc.session_name is None:
|
||||
return None
|
||||
filters["session_name"] = doc.session_name
|
||||
return filters
|
||||
|
||||
|
||||
async def query_external_vector_document_ids(
|
||||
workspace_name: str,
|
||||
observer: str,
|
||||
|
|
@ -473,6 +492,16 @@ def _dedup_key(
|
|||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _DocumentRowOp:
|
||||
kind: Literal["reinforce", "replace"]
|
||||
document_id: str
|
||||
incoming_times_derived: int = 1
|
||||
# When a reinforce skipped insert and the locked target is gone/deleted,
|
||||
# insert this document instead of dropping it.
|
||||
fallback_document: schemas.DocumentCreate | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class CreateDocumentsResult:
|
||||
created_documents: list[schemas.DocumentCreate] = field(default_factory=list)
|
||||
|
|
@ -515,6 +544,43 @@ async def create_documents(
|
|||
# Store (document_model, embedding) pairs - IDs aren't available until after commit
|
||||
docs_with_embeddings: list[tuple[models.Document, list[float]]] = []
|
||||
|
||||
# Resolve external-store dup candidates before the first DB statement.
|
||||
# None = pgvector in-place fallback; [] = skip semantic (no external I/O under db).
|
||||
semantic_candidates: list[list[str] | None] = [None] * len(documents)
|
||||
if deduplicate and not _uses_pgvector():
|
||||
resolve_sem = asyncio.Semaphore(_SEMANTIC_CANDIDATE_CONCURRENCY)
|
||||
|
||||
async def _resolve_candidates(index: int, doc: schemas.DocumentCreate) -> None:
|
||||
filters = _semantic_dup_filters(doc)
|
||||
if filters is None or not doc.embedding:
|
||||
semantic_candidates[index] = []
|
||||
return
|
||||
async with resolve_sem:
|
||||
try:
|
||||
ids = await query_external_vector_document_ids(
|
||||
workspace_name=workspace_name,
|
||||
observer=observer,
|
||||
observed=observed,
|
||||
embedding=doc.embedding,
|
||||
top_k=_SEMANTIC_DUP_TOP_K,
|
||||
max_distance=_SEMANTIC_DUP_MAX_DISTANCE,
|
||||
filters=filters,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"External semantic-candidate resolve failed for %s/%s/%s",
|
||||
workspace_name,
|
||||
observer,
|
||||
observed,
|
||||
)
|
||||
semantic_candidates[index] = []
|
||||
return
|
||||
semantic_candidates[index] = ids or []
|
||||
|
||||
await asyncio.gather(
|
||||
*(_resolve_candidates(i, doc) for i, doc in enumerate(documents))
|
||||
)
|
||||
|
||||
# exact-content dedup (independent of `deduplicate`): pre-fetch
|
||||
# existing live documents whose normalized content matches anything in this
|
||||
# batch, scoped to (workspace, observer, observed). The SQL normalization must
|
||||
|
|
@ -563,12 +629,14 @@ async def create_documents(
|
|||
# Tracks dedup keys already accepted from this batch so exact
|
||||
# duplicates within a single inference call collapse to one document.
|
||||
seen_in_batch: set[tuple[str, str, str | None]] = set()
|
||||
row_ops: list[_DocumentRowOp] = []
|
||||
pending_times_derived: dict[str, int] = {}
|
||||
|
||||
exact_dup_existing_count = 0
|
||||
exact_dup_in_batch_count = 0
|
||||
semantic_dup_rejected_count = 0
|
||||
semantic_dup_replaced_count = 0
|
||||
for doc in documents:
|
||||
for index, doc in enumerate(documents):
|
||||
try:
|
||||
# Session-purity invariant: an explicit document must always carry
|
||||
# the session it was derived from. Refuse to write session-less
|
||||
|
|
@ -598,88 +666,107 @@ async def create_documents(
|
|||
# the re-derivation as reinforcement on the existing row.
|
||||
existing_match = existing_by_key.get(dedup_key)
|
||||
if existing_match is not None:
|
||||
# Reinforce the existing row. greatest(...) keeps the bump atomic
|
||||
# server-side (concurrent workers can't lose an increment) while
|
||||
# still honoring an incoming doc that already carries accumulated
|
||||
# reinforcement (times_derived > 1, e.g. a future re-ingestion or
|
||||
# collection-merge path). Mirrors the superior-replacement branch
|
||||
# in is_rejected_duplicate.
|
||||
existing_match.times_derived = func.greatest(
|
||||
models.Document.times_derived + 1,
|
||||
doc.times_derived,
|
||||
current_td = pending_times_derived.get(
|
||||
existing_match.id, existing_match.times_derived
|
||||
)
|
||||
pending_times_derived[existing_match.id] = max(
|
||||
current_td + 1, doc.times_derived
|
||||
)
|
||||
row_ops.append(
|
||||
_DocumentRowOp(
|
||||
"reinforce",
|
||||
existing_match.id,
|
||||
doc.times_derived,
|
||||
fallback_document=doc,
|
||||
)
|
||||
)
|
||||
await db.flush()
|
||||
exact_dup_existing_count += 1
|
||||
continue
|
||||
|
||||
# for each document, if deduplicate is True, perform a process
|
||||
# that checks against existing documents and either rejects this document
|
||||
# as a duplicate OR deletes an existing document that is a duplicate.
|
||||
if deduplicate:
|
||||
duplicate_result = await is_rejected_duplicate(
|
||||
db, doc, workspace_name, observer=observer, observed=observed
|
||||
duplicate_result, existing_dup = await _semantic_dup_decision(
|
||||
db,
|
||||
doc,
|
||||
workspace_name,
|
||||
observer=observer,
|
||||
observed=observed,
|
||||
candidate_document_ids=semantic_candidates[index],
|
||||
)
|
||||
if duplicate_result is SemanticRejectionResult.REPLACED_EXISTING:
|
||||
# Existing doc was soft-deleted in favor of this one; the
|
||||
# new doc still gets inserted below.
|
||||
if (
|
||||
duplicate_result is SemanticRejectionResult.REPLACED_EXISTING
|
||||
and existing_dup is not None
|
||||
):
|
||||
current_td = pending_times_derived.get(
|
||||
existing_dup.id, existing_dup.times_derived
|
||||
)
|
||||
doc.times_derived = max(doc.times_derived, current_td + 1)
|
||||
pending_times_derived[existing_dup.id] = doc.times_derived
|
||||
row_ops.append(_DocumentRowOp("replace", existing_dup.id))
|
||||
semantic_dup_replaced_count += 1
|
||||
elif duplicate_result is SemanticRejectionResult.REJECTED:
|
||||
elif (
|
||||
duplicate_result is SemanticRejectionResult.REJECTED
|
||||
and existing_dup is not None
|
||||
):
|
||||
current_td = pending_times_derived.get(
|
||||
existing_dup.id, existing_dup.times_derived
|
||||
)
|
||||
pending_times_derived[existing_dup.id] = max(
|
||||
current_td + 1, doc.times_derived
|
||||
)
|
||||
row_ops.append(
|
||||
_DocumentRowOp(
|
||||
"reinforce",
|
||||
existing_dup.id,
|
||||
doc.times_derived,
|
||||
fallback_document=doc,
|
||||
)
|
||||
)
|
||||
semantic_dup_rejected_count += 1
|
||||
continue
|
||||
|
||||
metadata_dict = doc.metadata.model_dump(exclude_none=True)
|
||||
|
||||
# Determine if we need to persist embeddings to postgres
|
||||
# True when: TYPE=pgvector OR still migrating (dual-write to both stores)
|
||||
store_embeddings_in_postgres = (
|
||||
settings.VECTOR_STORE.TYPE == "pgvector"
|
||||
or not settings.VECTOR_STORE.MIGRATED
|
||||
new_doc = _document_model_from_create(
|
||||
doc, workspace_name=workspace_name, observer=observer, observed=observed
|
||||
)
|
||||
|
||||
if store_embeddings_in_postgres and doc.embedding:
|
||||
new_doc = models.Document(
|
||||
workspace_name=workspace_name,
|
||||
observer=observer,
|
||||
observed=observed,
|
||||
content=doc.content,
|
||||
level=doc.level,
|
||||
times_derived=doc.times_derived,
|
||||
internal_metadata=metadata_dict,
|
||||
session_name=doc.session_name,
|
||||
embedding=doc.embedding,
|
||||
# Tree linkage column
|
||||
source_ids=doc.source_ids,
|
||||
)
|
||||
else:
|
||||
new_doc = models.Document(
|
||||
workspace_name=workspace_name,
|
||||
observer=observer,
|
||||
observed=observed,
|
||||
content=doc.content,
|
||||
level=doc.level,
|
||||
times_derived=doc.times_derived,
|
||||
internal_metadata=metadata_dict,
|
||||
session_name=doc.session_name,
|
||||
# Tree linkage column
|
||||
source_ids=doc.source_ids,
|
||||
)
|
||||
|
||||
if doc.embedding:
|
||||
new_doc.sync_state = "pending"
|
||||
honcho_documents.append(new_doc)
|
||||
accepted_documents.append(doc)
|
||||
|
||||
# Track embedding for vector store (ID will be available after commit)
|
||||
if doc.embedding:
|
||||
docs_with_embeddings.append((new_doc, doc.embedding))
|
||||
|
||||
except IntegrityError as e:
|
||||
await db.rollback()
|
||||
raise ValidationException(
|
||||
"Failed to create documents due to integrity constraint violation"
|
||||
) from e
|
||||
except SQLAlchemyError:
|
||||
# Dead transaction: continuing would cascade PendingRollbackErrors.
|
||||
await db.rollback()
|
||||
raise
|
||||
except Exception as e:
|
||||
# Per-document failures (bad content, metadata, token overflow).
|
||||
logger.error(
|
||||
f"Error adding new document to {workspace_name}/{doc.session_name}/{observer}/{observed}: {e}"
|
||||
)
|
||||
continue
|
||||
|
||||
try:
|
||||
fallback_docs = await _apply_document_row_updates(
|
||||
db,
|
||||
row_ops,
|
||||
workspace_name=workspace_name,
|
||||
observer=observer,
|
||||
observed=observed,
|
||||
)
|
||||
for fallback_doc in fallback_docs:
|
||||
new_doc = _document_model_from_create(
|
||||
fallback_doc,
|
||||
workspace_name=workspace_name,
|
||||
observer=observer,
|
||||
observed=observed,
|
||||
)
|
||||
honcho_documents.append(new_doc)
|
||||
accepted_documents.append(fallback_doc)
|
||||
if fallback_doc.embedding:
|
||||
docs_with_embeddings.append((new_doc, fallback_doc.embedding))
|
||||
db.add_all(honcho_documents)
|
||||
# NOTE
|
||||
# If the process crashes after this commit but before vector upsert completes,
|
||||
|
|
@ -775,6 +862,11 @@ async def create_documents(
|
|||
raise ValidationException(
|
||||
"Failed to create documents due to integrity constraint violation"
|
||||
) from e
|
||||
except DBAPIError:
|
||||
# Leave the session clean for callers that own it (e.g. a deadlock
|
||||
# at the final commit); the queue layer classifies and retries.
|
||||
await db.rollback()
|
||||
raise
|
||||
|
||||
return CreateDocumentsResult(
|
||||
created_documents=accepted_documents,
|
||||
|
|
@ -1152,12 +1244,163 @@ async def create_observations(
|
|||
return honcho_documents
|
||||
|
||||
|
||||
def _document_model_from_create(
|
||||
doc: schemas.DocumentCreate,
|
||||
*,
|
||||
workspace_name: str,
|
||||
observer: str,
|
||||
observed: str,
|
||||
) -> models.Document:
|
||||
metadata_dict = doc.metadata.model_dump(exclude_none=True)
|
||||
store_embeddings_in_postgres = (
|
||||
settings.VECTOR_STORE.TYPE == "pgvector" or not settings.VECTOR_STORE.MIGRATED
|
||||
)
|
||||
if store_embeddings_in_postgres and doc.embedding:
|
||||
new_doc = models.Document(
|
||||
workspace_name=workspace_name,
|
||||
observer=observer,
|
||||
observed=observed,
|
||||
content=doc.content,
|
||||
level=doc.level,
|
||||
times_derived=doc.times_derived,
|
||||
internal_metadata=metadata_dict,
|
||||
session_name=doc.session_name,
|
||||
embedding=doc.embedding,
|
||||
source_ids=doc.source_ids,
|
||||
)
|
||||
else:
|
||||
new_doc = models.Document(
|
||||
workspace_name=workspace_name,
|
||||
observer=observer,
|
||||
observed=observed,
|
||||
content=doc.content,
|
||||
level=doc.level,
|
||||
times_derived=doc.times_derived,
|
||||
internal_metadata=metadata_dict,
|
||||
session_name=doc.session_name,
|
||||
source_ids=doc.source_ids,
|
||||
)
|
||||
if doc.embedding:
|
||||
new_doc.sync_state = "pending"
|
||||
return new_doc
|
||||
|
||||
|
||||
async def _apply_document_row_updates(
|
||||
db: AsyncSession,
|
||||
ops: list[_DocumentRowOp],
|
||||
*,
|
||||
workspace_name: str,
|
||||
observer: str,
|
||||
observed: str,
|
||||
) -> list[schemas.DocumentCreate]:
|
||||
"""Lock target rows by id, apply ops, return fallbacks for vanished targets."""
|
||||
if not ops:
|
||||
return []
|
||||
# Deadlock fix: lock in id order (IN-clause order is ignored).
|
||||
ids = sorted({op.document_id for op in ops})
|
||||
result = await db.execute(
|
||||
select(models.Document)
|
||||
.where(
|
||||
models.Document.id.in_(ids),
|
||||
models.Document.workspace_name == workspace_name,
|
||||
models.Document.observer == observer,
|
||||
models.Document.observed == observed,
|
||||
)
|
||||
.order_by(models.Document.id)
|
||||
.with_for_update()
|
||||
# Reload identity-map rows so the Python max() sees concurrent increments.
|
||||
.execution_options(populate_existing=True)
|
||||
)
|
||||
locked = {doc.id: doc for doc in result.scalars()}
|
||||
now = datetime.datetime.now(datetime.UTC)
|
||||
fallbacks: list[schemas.DocumentCreate] = []
|
||||
stale_at_lock = {
|
||||
op.document_id
|
||||
for op in ops
|
||||
if (locked_row := locked.get(op.document_id)) is None
|
||||
or locked_row.deleted_at is not None
|
||||
}
|
||||
for op in ops:
|
||||
row = locked.get(op.document_id)
|
||||
if op.kind == "replace":
|
||||
if row is not None and row.deleted_at is None:
|
||||
row.deleted_at = now
|
||||
continue
|
||||
# reinforce
|
||||
if op.document_id in stale_at_lock:
|
||||
if op.fallback_document is not None:
|
||||
fallbacks.append(op.fallback_document)
|
||||
continue
|
||||
if row is None or row.deleted_at is not None:
|
||||
# An earlier op in this batch replaced this row.
|
||||
continue
|
||||
row.times_derived = max(row.times_derived + 1, op.incoming_times_derived)
|
||||
await db.flush()
|
||||
return fallbacks
|
||||
|
||||
|
||||
class SemanticRejectionResult(Enum):
|
||||
NOT_DUPLICATE = 0
|
||||
REPLACED_EXISTING = 1
|
||||
REJECTED = 2
|
||||
|
||||
|
||||
async def _semantic_dup_decision(
|
||||
db: AsyncSession,
|
||||
doc: schemas.DocumentCreate,
|
||||
workspace_name: str,
|
||||
*,
|
||||
observer: str,
|
||||
observed: str,
|
||||
candidate_document_ids: list[str] | None = None,
|
||||
) -> tuple[SemanticRejectionResult, models.Document | None]:
|
||||
"""Classify a semantic duplicate without writing."""
|
||||
filters = _semantic_dup_filters(doc)
|
||||
if filters is None:
|
||||
return SemanticRejectionResult.NOT_DUPLICATE, None
|
||||
|
||||
if candidate_document_ids is not None:
|
||||
similar_docs: Sequence[models.Document] = await fetch_documents_by_ids(
|
||||
db=db,
|
||||
workspace_name=workspace_name,
|
||||
observer=observer,
|
||||
observed=observed,
|
||||
document_ids=candidate_document_ids,
|
||||
filters=filters,
|
||||
)
|
||||
elif _uses_pgvector():
|
||||
if not doc.embedding:
|
||||
# Match external-store path: never embed under an open session.
|
||||
return SemanticRejectionResult.NOT_DUPLICATE, None
|
||||
similar_docs = await query_documents(
|
||||
db=db,
|
||||
workspace_name=workspace_name,
|
||||
query=doc.content,
|
||||
observer=observer,
|
||||
observed=observed,
|
||||
filters=filters,
|
||||
max_distance=_SEMANTIC_DUP_MAX_DISTANCE,
|
||||
top_k=_SEMANTIC_DUP_TOP_K,
|
||||
embedding=doc.embedding,
|
||||
)
|
||||
else:
|
||||
return SemanticRejectionResult.NOT_DUPLICATE, None
|
||||
|
||||
if not similar_docs:
|
||||
return SemanticRejectionResult.NOT_DUPLICATE, None
|
||||
|
||||
existing_doc = similar_docs[0]
|
||||
tokens_new = set(embedding_client.encoding.encode(doc.content))
|
||||
tokens_existing = set(embedding_client.encoding.encode(existing_doc.content))
|
||||
unique_new = len(tokens_new - tokens_existing)
|
||||
unique_existing = len(tokens_existing - tokens_new)
|
||||
score_new = len(tokens_new) + (unique_new * 10)
|
||||
score_existing = len(tokens_existing) + (unique_existing * 10)
|
||||
if score_new >= score_existing:
|
||||
return SemanticRejectionResult.REPLACED_EXISTING, existing_doc
|
||||
return SemanticRejectionResult.REJECTED, existing_doc
|
||||
|
||||
|
||||
async def is_rejected_duplicate(
|
||||
db: AsyncSession,
|
||||
doc: schemas.DocumentCreate,
|
||||
|
|
@ -1165,90 +1408,29 @@ async def is_rejected_duplicate(
|
|||
*,
|
||||
observer: str,
|
||||
observed: str,
|
||||
candidate_document_ids: list[str] | None = None,
|
||||
) -> SemanticRejectionResult:
|
||||
"""
|
||||
Check if a document is a duplicate of an existing document.
|
||||
|
||||
Uses: 1) Cosine similarity (>=0.95), 2) Token diff for retention.
|
||||
|
||||
Returns True if both:
|
||||
- the document is deemed a duplicate of an existing document
|
||||
- the existing document is deemed a superior duplicate
|
||||
|
||||
If the document is not a duplicate, returns False.
|
||||
|
||||
If the document is a duplicate AND the new document is superior,
|
||||
deletes the existing document and returns False. In this case
|
||||
``doc.times_derived`` is updated in place to carry the replaced
|
||||
document's reinforcement count forward.
|
||||
|
||||
If the document is a duplicate AND the existing document is superior,
|
||||
increments the existing document's ``times_derived`` to record the
|
||||
reinforcement, then returns True.
|
||||
|
||||
Merges are scoped so they never cross document levels, and never cross
|
||||
sessions for explicit-level documents (session-purity invariant: an
|
||||
explicit document records what was derived from exactly one session, so
|
||||
a near-duplicate from another session must not reinforce or replace it).
|
||||
"""
|
||||
filters: dict[str, Any] = {"level": doc.level}
|
||||
if doc.level == "explicit":
|
||||
if doc.session_name is None:
|
||||
# create_documents refuses session-less explicit documents; if one
|
||||
# reaches here anyway it has no valid merge partner.
|
||||
return SemanticRejectionResult.NOT_DUPLICATE
|
||||
filters["session_name"] = doc.session_name
|
||||
|
||||
# Step 1: Find potential duplicates using cosine similarity
|
||||
similar_docs = await query_documents(
|
||||
db=db,
|
||||
workspace_name=workspace_name,
|
||||
query=doc.content,
|
||||
"""Classify a semantic duplicate and apply the corresponding row write."""
|
||||
result, existing_doc = await _semantic_dup_decision(
|
||||
db,
|
||||
doc,
|
||||
workspace_name,
|
||||
observer=observer,
|
||||
observed=observed,
|
||||
filters=filters,
|
||||
max_distance=0.05,
|
||||
top_k=1,
|
||||
embedding=doc.embedding,
|
||||
candidate_document_ids=candidate_document_ids,
|
||||
)
|
||||
|
||||
if not similar_docs:
|
||||
return SemanticRejectionResult.NOT_DUPLICATE
|
||||
|
||||
existing_doc = similar_docs[0]
|
||||
|
||||
# Step 2: Determine which has more information using token set difference
|
||||
tokens_new = set(embedding_client.encoding.encode(doc.content))
|
||||
tokens_existing = set(embedding_client.encoding.encode(existing_doc.content))
|
||||
|
||||
unique_new = len(tokens_new - tokens_existing)
|
||||
unique_existing = len(tokens_existing - tokens_new)
|
||||
|
||||
score_new = len(tokens_new) + (unique_new * 10)
|
||||
score_existing = len(tokens_existing) + (unique_existing * 10)
|
||||
|
||||
# If new document has more or equal information, keep it and delete existing
|
||||
if score_new >= score_existing:
|
||||
if existing_doc is None:
|
||||
return result
|
||||
if result is SemanticRejectionResult.REPLACED_EXISTING:
|
||||
logger.debug(
|
||||
"[DUPLICATE DETECTION] Deleting existing in favor of new. new=%r, existing=%r.",
|
||||
doc.content,
|
||||
existing_doc.content,
|
||||
)
|
||||
# Carry the reinforcement count forward so replacing a duplicate counts as
|
||||
# another derivation rather than resetting times_derived to 1.
|
||||
doc.times_derived = max(doc.times_derived, existing_doc.times_derived + 1)
|
||||
# Soft-delete the existing document - reconciliation will clean up vectors and hard-delete
|
||||
existing_doc.deleted_at = datetime.datetime.now(datetime.timezone.utc)
|
||||
existing_doc.deleted_at = datetime.datetime.now(datetime.UTC)
|
||||
await db.flush()
|
||||
return (
|
||||
SemanticRejectionResult.REPLACED_EXISTING
|
||||
) # Don't reject the new document
|
||||
|
||||
# Existing document has more information, reject the new one but record the
|
||||
# reinforcement: a semantic duplicate was derived again. greatest(...) keeps
|
||||
# the increment atomic server-side -- concurrent workers reinforcing the same
|
||||
# document must not lose updates -- while still honoring an incoming doc that
|
||||
# already carries accumulated reinforcement (times_derived > 1).
|
||||
return result
|
||||
existing_doc.times_derived = func.greatest(
|
||||
models.Document.times_derived + 1,
|
||||
doc.times_derived,
|
||||
|
|
@ -1259,7 +1441,7 @@ async def is_rejected_duplicate(
|
|||
doc.content,
|
||||
existing_doc.content,
|
||||
)
|
||||
return SemanticRejectionResult.REJECTED
|
||||
return result
|
||||
|
||||
|
||||
async def cleanup_soft_deleted_documents(
|
||||
|
|
@ -1284,7 +1466,7 @@ async def cleanup_soft_deleted_documents(
|
|||
Returns:
|
||||
Count of documents cleaned up (only those where vector deletion succeeded).
|
||||
"""
|
||||
cutoff = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(
|
||||
cutoff = datetime.datetime.now(datetime.UTC) - datetime.timedelta(
|
||||
minutes=older_than_minutes
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -25,6 +25,7 @@ from src.utils.representation import (
|
|||
Representation,
|
||||
allowlist_safe_levels,
|
||||
)
|
||||
from src.utils.sanitization import strip_nul
|
||||
from src.utils.types import embedding_call_purpose
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
|
@ -38,10 +39,21 @@ def _observation_text(obs: ExplicitObservation | DeductiveObservation) -> str:
|
|||
def _normalized_observation(
|
||||
obs: ExplicitObservation | DeductiveObservation,
|
||||
) -> ExplicitObservation | DeductiveObservation:
|
||||
"""Return an observation with its persisted/embed text normalized."""
|
||||
text = _observation_text(obs).strip()
|
||||
"""Return an observation with its persisted/embed text normalized.
|
||||
|
||||
NUL bytes are removed here rather than closer to the database so that the
|
||||
text that gets embedded is the same text that gets stored.
|
||||
"""
|
||||
text = strip_nul(_observation_text(obs)).strip()
|
||||
if isinstance(obs, DeductiveObservation):
|
||||
return obs.model_copy(update={"conclusion": text})
|
||||
return obs.model_copy(
|
||||
update={
|
||||
"conclusion": text,
|
||||
# Premises ride along in internal_metadata, and jsonb rejects
|
||||
# NUL in strings just as text columns do.
|
||||
"premises": strip_nul(obs.premises),
|
||||
}
|
||||
)
|
||||
return obs.model_copy(update={"content": text})
|
||||
|
||||
|
||||
|
|
@ -87,10 +99,15 @@ class RepresentationManager:
|
|||
logger.debug("No observations to save")
|
||||
return empty_result
|
||||
|
||||
# Normalize before the emptiness check: str.strip() does not remove
|
||||
# NUL, so content that normalizes away has to be dropped afterwards.
|
||||
all_observations = [
|
||||
_normalized_observation(obs)
|
||||
for obs in representation.deductive + representation.explicit
|
||||
if _observation_text(obs).strip()
|
||||
normalized
|
||||
for normalized in (
|
||||
_normalized_observation(obs)
|
||||
for obs in representation.deductive + representation.explicit
|
||||
)
|
||||
if _observation_text(normalized)
|
||||
]
|
||||
if not all_observations:
|
||||
logger.debug("No non-empty observations to save")
|
||||
|
|
|
|||
|
|
@ -27,6 +27,7 @@ from src.telemetry.events import (
|
|||
from src.telemetry.logging import log_performance_metrics
|
||||
from src.utils import summarizer
|
||||
from src.utils.queue_payload import (
|
||||
RETRY_ATTEMPTS_PAYLOAD_KEY,
|
||||
DeletionPayload,
|
||||
DreamPayload,
|
||||
ReconcilerPayload,
|
||||
|
|
@ -44,7 +45,11 @@ logging.getLogger("sqlalchemy.engine.Engine").disabled = True
|
|||
async def process_item(queue_item: models.QueueItem) -> None:
|
||||
"""Process a single item from the queue."""
|
||||
task_type = queue_item.task_type
|
||||
queue_payload = queue_item.payload
|
||||
# Drop the work-unit retry counter before payload validation: every payload
|
||||
# model sets extra="forbid", so leaving it in burns the item as
|
||||
# extra_forbidden on the reclaim that was supposed to retry it.
|
||||
queue_payload = dict(queue_item.payload or {})
|
||||
queue_payload.pop(RETRY_ATTEMPTS_PAYLOAD_KEY, None)
|
||||
workspace_name = queue_item.workspace_name
|
||||
|
||||
# Handle reconciler first - it's the only task type that doesn't require workspace_name
|
||||
|
|
|
|||
|
|
@ -25,6 +25,7 @@ from src.telemetry.sentry import with_sentry_transaction
|
|||
from src.utils.config_helpers import get_configuration
|
||||
from src.utils.formatting import format_new_turn_with_timestamp
|
||||
from src.utils.representation import PromptRepresentation, Representation
|
||||
from src.utils.retryable_errors import is_retryable_error
|
||||
from src.utils.tokens import track_deriver_input_tokens
|
||||
|
||||
from .prompts import estimate_deriver_prompt_tokens, minimal_deriver_prompt
|
||||
|
|
@ -344,6 +345,12 @@ async def process_representation_tasks_batch(
|
|||
)
|
||||
)
|
||||
|
||||
retryable = next(
|
||||
(exc for _, exc in save_errors if is_retryable_error(exc)),
|
||||
None,
|
||||
)
|
||||
if retryable is not None:
|
||||
raise retryable
|
||||
if save_errors and successful_observer_count == 0:
|
||||
details = "; ".join(
|
||||
f"{observer}: {exc.__class__.__name__}: {exc}"
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ import time
|
|||
from asyncio import Task
|
||||
from collections.abc import Iterable, Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from logging import getLogger
|
||||
from typing import Any, NamedTuple, cast
|
||||
|
||||
|
|
@ -15,7 +15,7 @@ from dotenv import load_dotenv
|
|||
from nanoid import generate as generate_nanoid
|
||||
from sentry_sdk.integrations.asyncio import AsyncioIntegration
|
||||
from sentry_sdk.integrations.sqlalchemy import SqlalchemyIntegration
|
||||
from sqlalchemy import and_, delete, or_, select, update
|
||||
from sqlalchemy import Text, and_, delete, literal, select, update
|
||||
from sqlalchemy.dialects.postgresql import insert
|
||||
from sqlalchemy.engine import CursorResult
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
|
@ -24,6 +24,11 @@ from sqlalchemy.sql import func
|
|||
from src import models
|
||||
from src.cache.client import close_cache, init_cache
|
||||
from src.config import settings
|
||||
from src.crud.deriver import (
|
||||
REPRESENTATION_WORK_UNIT_PREFIX,
|
||||
representation_batch_threshold_clause,
|
||||
unclaimed_work_unit_clause,
|
||||
)
|
||||
from src.dependencies import tracked_db
|
||||
from src.deriver.consumer import (
|
||||
process_item,
|
||||
|
|
@ -43,6 +48,8 @@ from src.reconciler import (
|
|||
from src.schemas import ResolvedConfiguration
|
||||
from src.telemetry import prometheus_metrics
|
||||
from src.telemetry.sentry import initialize_sentry
|
||||
from src.utils.queue_payload import RETRY_ATTEMPTS_PAYLOAD_KEY
|
||||
from src.utils.retryable_errors import is_retryable_error
|
||||
from src.utils.work_unit import parse_work_unit_key
|
||||
from src.webhooks.events import (
|
||||
QueueEmptyEvent,
|
||||
|
|
@ -53,6 +60,12 @@ logger = getLogger(__name__)
|
|||
|
||||
load_dotenv(override=True)
|
||||
|
||||
# Total processing attempts per work unit for transient errors. Count is
|
||||
# stored on the oldest unprocessed queue item so every deriver instance
|
||||
# shares one budget.
|
||||
MAX_RETRYABLE_ATTEMPTS = 3
|
||||
RETRY_BACKOFF_SECONDS = 1.0
|
||||
|
||||
|
||||
class WorkerOwnership(NamedTuple):
|
||||
"""Represents the instance of a work unit that a worker is processing."""
|
||||
|
|
@ -301,7 +314,7 @@ class QueueManager:
|
|||
async def cleanup_stale_work_units(self) -> None:
|
||||
"""Clean up stale work units"""
|
||||
async with tracked_db("cleanup_stale_work_units") as db:
|
||||
cutoff = datetime.now(timezone.utc) - timedelta(
|
||||
cutoff = datetime.now(UTC) - timedelta(
|
||||
minutes=settings.DERIVER.STALE_SESSION_TIMEOUT_MINUTES
|
||||
)
|
||||
|
||||
|
|
@ -345,7 +358,7 @@ class QueueManager:
|
|||
)
|
||||
|
||||
async with tracked_db("get_available_work_units") as db:
|
||||
representation_prefix = "representation:"
|
||||
representation_prefix = REPRESENTATION_WORK_UNIT_PREFIX
|
||||
token_stats_subq = (
|
||||
select(
|
||||
models.QueueItem.work_unit_key,
|
||||
|
|
@ -382,14 +395,7 @@ class QueueManager:
|
|||
token_stats_subq,
|
||||
work_units_subq.c.work_unit_key == token_stats_subq.c.work_unit_key,
|
||||
)
|
||||
.where(
|
||||
~select(models.ActiveQueueSession.id)
|
||||
.where(
|
||||
models.ActiveQueueSession.work_unit_key
|
||||
== work_units_subq.c.work_unit_key
|
||||
)
|
||||
.exists()
|
||||
)
|
||||
.where(unclaimed_work_unit_clause(work_units_subq.c.work_unit_key))
|
||||
.order_by(
|
||||
work_units_subq.c.oldest_created_at.asc(),
|
||||
work_units_subq.c.work_unit_key.asc(),
|
||||
|
|
@ -398,26 +404,13 @@ class QueueManager:
|
|||
)
|
||||
|
||||
# Apply batch threshold filter (skip if FLUSH_ENABLED is True)
|
||||
if not settings.DERIVER.FLUSH_ENABLED and work_unit_target_tokens > 0:
|
||||
max_age_seconds = settings.DERIVER.REPRESENTATION_BATCH_MAX_AGE_SECONDS
|
||||
threshold_clause = (
|
||||
func.coalesce(token_stats_subq.c.total_tokens, 0)
|
||||
>= work_unit_target_tokens
|
||||
)
|
||||
if max_age_seconds > 0:
|
||||
threshold_clause = or_(
|
||||
threshold_clause,
|
||||
token_stats_subq.c.oldest_created_at
|
||||
<= func.now() - timedelta(seconds=max_age_seconds),
|
||||
)
|
||||
query = query.where(
|
||||
or_(
|
||||
~work_units_subq.c.work_unit_key.startswith(
|
||||
representation_prefix
|
||||
),
|
||||
threshold_clause,
|
||||
)
|
||||
)
|
||||
threshold_clause = representation_batch_threshold_clause(
|
||||
work_unit_key=work_units_subq.c.work_unit_key,
|
||||
total_tokens=token_stats_subq.c.total_tokens,
|
||||
oldest_created_at=token_stats_subq.c.oldest_created_at,
|
||||
)
|
||||
if threshold_clause is not None:
|
||||
query = query.where(threshold_clause)
|
||||
|
||||
result = await db.execute(query)
|
||||
available_rows = result.all()
|
||||
|
|
@ -591,11 +584,24 @@ class QueueManager:
|
|||
items: list[QueueItem],
|
||||
work_unit_key: str,
|
||||
context: str,
|
||||
) -> None:
|
||||
) -> bool:
|
||||
"""
|
||||
Handle processing errors by marking queue items as errored, logging, and forwarding to Sentry.
|
||||
We only mark the first queue item as errored so we don't potentially throw away a batch. This allows us
|
||||
to incrementally attempt to process the batch while still maintaining progress in a work unit.
|
||||
Handle a processing error. Returns True when the caller should stop
|
||||
processing and release the work unit for a later re-claim.
|
||||
|
||||
Transient errors (is_retryable_error) get up to MAX_RETRYABLE_ATTEMPTS
|
||||
attempts per work unit: items stay unprocessed with no error recorded.
|
||||
The attempt count lives on the oldest unprocessed queue item so a
|
||||
different deriver instance continues the same budget after reclaim.
|
||||
Reprocessing is at-least-once, not idempotent: the batch is re-derived
|
||||
by a fresh LLM call, so identical text collapses via exact dedup and
|
||||
near-identical text via semantic dedup. Retries can therefore inflate
|
||||
times_derived and double-count LLM telemetry -- acceptable because the
|
||||
alternative is dropping the batch.
|
||||
|
||||
Terminal errors mark only the first queue item as errored so we don't
|
||||
potentially throw away a batch. This allows us to incrementally attempt
|
||||
to process the batch while still maintaining progress in a work unit.
|
||||
|
||||
Args:
|
||||
error: The exception that occurred
|
||||
|
|
@ -603,12 +609,37 @@ class QueueManager:
|
|||
work_unit_key: The work unit key for the queue items
|
||||
context: Context string describing what was being processed (e.g., "processing representation batch")
|
||||
"""
|
||||
if is_retryable_error(error):
|
||||
try:
|
||||
attempts = await self._get_work_unit_retry_attempts(work_unit_key) + 1
|
||||
if attempts < MAX_RETRYABLE_ATTEMPTS:
|
||||
await self._set_work_unit_retry_attempts(work_unit_key, attempts)
|
||||
logger.warning(
|
||||
"Transient error %s for work unit %s (attempt %d/%d); leaving items unprocessed for retry",
|
||||
context,
|
||||
work_unit_key,
|
||||
attempts,
|
||||
MAX_RETRYABLE_ATTEMPTS,
|
||||
exc_info=error,
|
||||
)
|
||||
return True
|
||||
except Exception: # noqa: BLE001
|
||||
logger.exception(
|
||||
"Retry-counter I/O failed for work unit %s; releasing %s without recording an attempt",
|
||||
work_unit_key,
|
||||
context,
|
||||
)
|
||||
return True
|
||||
|
||||
error_msg = f"{error.__class__.__name__}: {str(error)}"
|
||||
try:
|
||||
if items:
|
||||
# Clear retry metadata only after the terminal mark commits so a
|
||||
# failed mark leaves the shared budget intact for the next claim.
|
||||
await self.mark_queue_item_as_errored(
|
||||
items[0], work_unit_key, error_msg
|
||||
)
|
||||
await self._clear_work_unit_retry_attempts(work_unit_key)
|
||||
except Exception as mark_error:
|
||||
logger.error(
|
||||
f"Failed to mark queue items as errored for work unit {work_unit_key}: {mark_error}",
|
||||
|
|
@ -621,6 +652,7 @@ class QueueManager:
|
|||
)
|
||||
if settings.SENTRY.ENABLED:
|
||||
sentry_sdk.capture_exception(error)
|
||||
return False
|
||||
|
||||
async def process_work_unit(self, work_unit_key: str, worker_id: str) -> None:
|
||||
"""Process all queue items for a specific work unit by routing to the correct handler."""
|
||||
|
|
@ -686,12 +718,18 @@ class QueueManager:
|
|||
)
|
||||
queue_item_count += len(items_to_process)
|
||||
except Exception as e:
|
||||
await self._handle_processing_error(
|
||||
if await self._handle_processing_error(
|
||||
e,
|
||||
items_to_process,
|
||||
work_unit_key,
|
||||
f"processing {work_unit.task_type} batch",
|
||||
)
|
||||
):
|
||||
# Release the work unit (via the finally
|
||||
# below) and let a later poll re-claim it.
|
||||
await asyncio.sleep(
|
||||
self._jitter(RETRY_BACKOFF_SECONDS)
|
||||
)
|
||||
break
|
||||
|
||||
else:
|
||||
queue_item = await self.get_next_queue_item(
|
||||
|
|
@ -710,12 +748,16 @@ class QueueManager:
|
|||
)
|
||||
queue_item_count += 1
|
||||
except Exception as e:
|
||||
await self._handle_processing_error(
|
||||
if await self._handle_processing_error(
|
||||
e,
|
||||
[queue_item],
|
||||
work_unit_key,
|
||||
"processing queue item",
|
||||
)
|
||||
):
|
||||
await asyncio.sleep(
|
||||
self._jitter(RETRY_BACKOFF_SECONDS)
|
||||
)
|
||||
break
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
|
|
@ -1068,6 +1110,87 @@ class QueueManager:
|
|||
batch_max_tokens=batch_max_tokens,
|
||||
)
|
||||
|
||||
async def _oldest_unprocessed_item(
|
||||
self,
|
||||
db: AsyncSession,
|
||||
work_unit_key: str,
|
||||
*,
|
||||
for_update: bool = False,
|
||||
) -> models.QueueItem | None:
|
||||
stmt = (
|
||||
select(models.QueueItem)
|
||||
.where(
|
||||
models.QueueItem.work_unit_key == work_unit_key,
|
||||
models.QueueItem.processed.is_(False),
|
||||
)
|
||||
.order_by(models.QueueItem.id)
|
||||
.limit(1)
|
||||
)
|
||||
if for_update:
|
||||
stmt = stmt.with_for_update()
|
||||
result = await db.execute(stmt)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
async def _get_work_unit_retry_attempts(self, work_unit_key: str) -> int:
|
||||
"""Read the shared transient-failure attempt count for a work unit."""
|
||||
async with tracked_db("get_work_unit_retry_attempts") as db:
|
||||
item = await self._oldest_unprocessed_item(db, work_unit_key)
|
||||
if item is None:
|
||||
return 0
|
||||
raw = (item.payload or {}).get(RETRY_ATTEMPTS_PAYLOAD_KEY, 0)
|
||||
try:
|
||||
return max(0, int(raw))
|
||||
except (TypeError, ValueError):
|
||||
return 0
|
||||
|
||||
async def _set_work_unit_retry_attempts(
|
||||
self, work_unit_key: str, attempts: int
|
||||
) -> None:
|
||||
"""Persist the shared attempt count on the oldest unprocessed item."""
|
||||
async with tracked_db("set_work_unit_retry_attempts") as db:
|
||||
item = await self._oldest_unprocessed_item(
|
||||
db, work_unit_key, for_update=True
|
||||
)
|
||||
if item is None:
|
||||
await db.commit()
|
||||
return
|
||||
new_payload = dict(item.payload or {})
|
||||
new_payload[RETRY_ATTEMPTS_PAYLOAD_KEY] = attempts
|
||||
await db.execute(
|
||||
update(models.QueueItem)
|
||||
.where(models.QueueItem.id == item.id)
|
||||
.values(payload=new_payload)
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
async def _clear_work_unit_retry_attempts(self, work_unit_key: str) -> None:
|
||||
"""Drop the shared attempt count from remaining unprocessed items.
|
||||
|
||||
One statement on purpose: a multi-row ``SELECT ... FOR UPDATE`` here
|
||||
would take locks on ``queue`` in scan order, which is a deadlock partner
|
||||
for any other multi-row writer on the same table. The JSONB ``-``
|
||||
operator does the strip server-side, so no rows are locked ahead of the
|
||||
write and there is no lock order to get wrong.
|
||||
"""
|
||||
async with tracked_db("clear_work_unit_retry_attempts") as db:
|
||||
await db.execute(
|
||||
update(models.QueueItem)
|
||||
.where(
|
||||
models.QueueItem.work_unit_key == work_unit_key,
|
||||
models.QueueItem.processed.is_(False),
|
||||
models.QueueItem.payload.has_key(RETRY_ATTEMPTS_PAYLOAD_KEY),
|
||||
)
|
||||
.values(
|
||||
# literal(..., Text) is required: an untyped bind leaves
|
||||
# Postgres unable to pick between jsonb - text and its
|
||||
# integer/array siblings.
|
||||
payload=models.QueueItem.payload.op("-")(
|
||||
literal(RETRY_ATTEMPTS_PAYLOAD_KEY, Text)
|
||||
)
|
||||
)
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
async def mark_queue_items_as_processed(
|
||||
self, items: list[QueueItem], work_unit_key: str
|
||||
) -> None:
|
||||
|
|
|
|||
|
|
@ -30,12 +30,12 @@ from typing import Any
|
|||
|
||||
from sqlalchemy import select, update
|
||||
from sqlalchemy.dialects.postgresql import array
|
||||
from sqlalchemy.orm import load_only
|
||||
from sqlalchemy.sql.functions import func
|
||||
|
||||
from src import crud, models
|
||||
from src.config import settings
|
||||
from src.crud.scope import ScopeBackfillState
|
||||
from src.crud.session import is_peer_in_session
|
||||
from src.dependencies import tracked_db
|
||||
from src.embedding_client import embedding_client
|
||||
from src.schemas import DreamType
|
||||
|
|
@ -49,6 +49,10 @@ logger = logging.getLogger(__name__)
|
|||
# was copied from. The presence of this key is the idempotency marker.
|
||||
COPIED_FROM_KEY = "copied_from"
|
||||
|
||||
# Specs embedded, written, and synced per pass. Bounds the live embeddings
|
||||
# (~40KB each as Python floats) so a large session cannot OOM the deriver.
|
||||
BACKFILL_CHUNK_SIZE = 500
|
||||
|
||||
|
||||
def _store_embeddings_in_postgres() -> bool:
|
||||
"""Whether document embeddings are persisted to the postgres column.
|
||||
|
|
@ -173,7 +177,23 @@ async def _run_backfill(
|
|||
plans: list[_CopySpec] = []
|
||||
async with tracked_db("scope_backfill.plan") as db:
|
||||
source_result = await db.execute(
|
||||
select(models.Document).where(
|
||||
select(models.Document)
|
||||
.options(
|
||||
load_only(
|
||||
models.Document.id,
|
||||
models.Document.workspace_name,
|
||||
models.Document.observer,
|
||||
models.Document.observed,
|
||||
models.Document.content,
|
||||
models.Document.level,
|
||||
models.Document.times_derived,
|
||||
models.Document.internal_metadata,
|
||||
models.Document.session_name,
|
||||
models.Document.source_ids,
|
||||
models.Document.deleted_at,
|
||||
)
|
||||
)
|
||||
.where(
|
||||
models.Document.workspace_name == workspace_name,
|
||||
models.Document.session_name == session_name,
|
||||
models.Document.level == "explicit",
|
||||
|
|
@ -196,7 +216,19 @@ async def _run_backfill(
|
|||
# by (observed, copied_from). Includes soft-deleted rows: those are
|
||||
# restore candidates, not blockers.
|
||||
copies_result = await db.execute(
|
||||
select(models.Document).where(
|
||||
select(models.Document)
|
||||
.options(
|
||||
load_only(
|
||||
models.Document.id,
|
||||
models.Document.workspace_name,
|
||||
models.Document.observer,
|
||||
models.Document.observed,
|
||||
models.Document.session_name,
|
||||
models.Document.internal_metadata,
|
||||
models.Document.deleted_at,
|
||||
)
|
||||
)
|
||||
.where(
|
||||
models.Document.workspace_name == workspace_name,
|
||||
models.Document.observer == scope_peer,
|
||||
models.Document.session_name == session_name,
|
||||
|
|
@ -216,12 +248,13 @@ async def _run_backfill(
|
|||
key = (source.observed, source.id)
|
||||
if key in live_copies:
|
||||
continue
|
||||
# Vectors hydrate per chunk; plans only carry ids + content.
|
||||
plans.append(
|
||||
_CopySpec(
|
||||
observed=source.observed,
|
||||
source_id=source.id,
|
||||
content=source.content,
|
||||
embedding=_embedding_as_list(source.embedding),
|
||||
embedding=None,
|
||||
internal_metadata=dict(source.internal_metadata),
|
||||
times_derived=source.times_derived,
|
||||
source_ids=list(source.source_ids)
|
||||
|
|
@ -235,7 +268,54 @@ async def _run_backfill(
|
|||
if not plans:
|
||||
return 0, set()
|
||||
|
||||
# Phase 2 (no DB): fill missing embeddings. Source rows have NULL
|
||||
# Phases 2-4 run per chunk so only one chunk's embeddings are alive at a
|
||||
# time; each chunk's vectors are dropped once synced.
|
||||
store_in_postgres = _store_embeddings_in_postgres()
|
||||
touched_observed: set[str] = set()
|
||||
copied = 0
|
||||
for start in range(0, len(plans), BACKFILL_CHUNK_SIZE):
|
||||
chunk = plans[start : start + BACKFILL_CHUNK_SIZE]
|
||||
if not await _copy_chunk(
|
||||
workspace_name, scope_peer, session_name, chunk, store_in_postgres
|
||||
):
|
||||
return None
|
||||
copied += len(chunk)
|
||||
touched_observed.update(spec.observed for spec in chunk)
|
||||
for spec in chunk:
|
||||
spec.embedding = None
|
||||
|
||||
return copied, touched_observed
|
||||
|
||||
|
||||
async def _hydrate_chunk_embeddings(
|
||||
workspace_name: str, plans: list[_CopySpec]
|
||||
) -> None:
|
||||
"""Load this chunk's source embeddings from postgres (if any)."""
|
||||
source_ids = [spec.source_id for spec in plans]
|
||||
async with tracked_db("scope_backfill.hydrate_embeddings") as db:
|
||||
result = await db.execute(
|
||||
select(models.Document.id, models.Document.embedding).where(
|
||||
models.Document.workspace_name == workspace_name,
|
||||
models.Document.id.in_(source_ids),
|
||||
)
|
||||
)
|
||||
by_id = {row.id: _embedding_as_list(row.embedding) for row in result.all()}
|
||||
for spec in plans:
|
||||
spec.embedding = by_id.get(spec.source_id)
|
||||
|
||||
|
||||
async def _copy_chunk(
|
||||
workspace_name: str,
|
||||
scope_peer: str,
|
||||
session_name: str,
|
||||
plans: list[_CopySpec],
|
||||
store_in_postgres: bool,
|
||||
) -> bool:
|
||||
"""Embed, write, and sync one chunk. False if the session left the scope."""
|
||||
# Phase 2a (DB): pull this chunk's embeddings only.
|
||||
await _hydrate_chunk_embeddings(workspace_name, plans)
|
||||
|
||||
# Phase 2b (no DB): fill missing embeddings. Source rows have NULL
|
||||
# embeddings on external-store deployments (and soft-deleted copies may
|
||||
# have lost their vectors) — re-embed via the embedding API only; no LLM.
|
||||
missing = [spec for spec in plans if spec.embedding is None]
|
||||
|
|
@ -254,17 +334,22 @@ async def _run_backfill(
|
|||
spec.embedding = embedding
|
||||
|
||||
# Phase 3 (DB): write the copies.
|
||||
store_in_postgres = _store_embeddings_in_postgres()
|
||||
touched_observed = {spec.observed for spec in plans}
|
||||
new_rows: list[models.Document] = []
|
||||
async with tracked_db("scope_backfill.write") as db:
|
||||
# scope_backfill and scope_removal carry different work-unit keys, so
|
||||
# nothing orders them: a removal enqueued right after the add (or one
|
||||
# that landed while phase 2 was embedding) can sweep the scope before
|
||||
# these copies exist. Re-checking membership here, in the transaction
|
||||
# that inserts, keeps a removed session from being copied back in.
|
||||
if not await is_peer_in_session(db, workspace_name, session_name, scope_peer):
|
||||
return None
|
||||
# Row-lock active membership for this txn so a concurrent leave
|
||||
# (``left_at``) cannot commit between the check and the inserts.
|
||||
membership = await db.scalar(
|
||||
select(models.SessionPeer.peer_name)
|
||||
.where(models.SessionPeer.workspace_name == workspace_name)
|
||||
.where(models.SessionPeer.session_name == session_name)
|
||||
.where(models.SessionPeer.peer_name == scope_peer)
|
||||
.where(models.SessionPeer.left_at.is_(None))
|
||||
.with_for_update()
|
||||
.limit(1)
|
||||
)
|
||||
if membership is None:
|
||||
return False
|
||||
|
||||
for observed in sorted(touched_observed):
|
||||
await crud.get_or_create_collection(
|
||||
|
|
@ -319,8 +404,7 @@ async def _run_backfill(
|
|||
# Phase 4: sync to the external vector store (or mark synced in pgvector
|
||||
# mode). Failures leave rows in sync_state='pending' for the reconciler.
|
||||
await _sync_copies_to_vector_store(workspace_name, scope_peer, plans, copied_ids)
|
||||
|
||||
return len(plans), touched_observed
|
||||
return True
|
||||
|
||||
|
||||
async def _sync_copies_to_vector_store(
|
||||
|
|
|
|||
|
|
@ -14,7 +14,12 @@ from nanoid import generate as generate_nanoid
|
|||
from pydantic import BaseModel
|
||||
|
||||
from src import crud
|
||||
from src.config import ConfiguredModelSettings, ReasoningLevel, settings
|
||||
from src.config import (
|
||||
ConfiguredModelSettings,
|
||||
DialecticLevelSettings,
|
||||
ReasoningLevel,
|
||||
settings,
|
||||
)
|
||||
from src.dependencies import tracked_db
|
||||
from src.dialectic import prompts
|
||||
from src.embedding_client import embedding_client
|
||||
|
|
@ -139,6 +144,20 @@ class DialecticAgent:
|
|||
tools = [t for t in tools if t.get("name") != "get_reasoning_chain"]
|
||||
return tools
|
||||
|
||||
def _tool_choice(
|
||||
self, level_settings: DialecticLevelSettings
|
||||
) -> str | dict[str, Any] | None:
|
||||
"""Pick the tool_choice for this query.
|
||||
|
||||
Defaults to whatever the reasoning level configures. Subclasses override
|
||||
when the agent has no prefetched corpus to fall back on and so must
|
||||
search before it can answer. Forcing "required"/"any" here costs exactly
|
||||
one tool round rather than pinning the loop: `execute_tool_loop` relaxes
|
||||
it to "auto" after the first iteration so the model can still stop and
|
||||
synthesize.
|
||||
"""
|
||||
return level_settings.TOOL_CHOICE
|
||||
|
||||
async def _initialize_session_history(self) -> None:
|
||||
"""Fetch and inject session history into the system prompt if configured."""
|
||||
if self._session_history_initialized:
|
||||
|
|
@ -505,7 +524,7 @@ class DialecticAgent:
|
|||
prompt="", # Ignored since we pass messages
|
||||
max_tokens=max_tokens,
|
||||
tools=tools,
|
||||
tool_choice=level_settings.TOOL_CHOICE,
|
||||
tool_choice=self._tool_choice(level_settings),
|
||||
tool_executor=tool_executor,
|
||||
max_tool_iterations=level_settings.MAX_TOOL_ITERATIONS,
|
||||
messages=self.messages,
|
||||
|
|
@ -581,7 +600,7 @@ class DialecticAgent:
|
|||
stream=True,
|
||||
stream_final_only=True,
|
||||
tools=tools,
|
||||
tool_choice=level_settings.TOOL_CHOICE,
|
||||
tool_choice=self._tool_choice(level_settings),
|
||||
tool_executor=tool_executor,
|
||||
max_tool_iterations=level_settings.MAX_TOOL_ITERATIONS,
|
||||
messages=self.messages,
|
||||
|
|
|
|||
|
|
@ -396,7 +396,11 @@ If this query is restricted to a session or a set of sessions, message tools alr
|
|||
|
||||
4. **Attribute**. Every fact you state names the peer it is about. If it is a cross-peer view, also name whose model it came from. Example: "Alice is a violinist." / "From Bob's model of Alice, …"
|
||||
|
||||
5. **Synthesize**. Answer the question. Quote exact names, dates, and numbers. For aggregations, list findings per peer. Do not narrate tool use.
|
||||
5. **Synthesize**. Answer the question. Quote exact names, dates, and numbers. For aggregations, list findings per peer. Do not narrate tool use, and do not describe a search you did not run.
|
||||
|
||||
## NO CLARIFYING QUESTIONS
|
||||
|
||||
Your answer goes to a program, not to someone who can reply. No one will answer a question you ask, approve a plan you propose, or pick from options you offer — your response ends the exchange. So never ask which lookup to run, never lay out a plan and stop, never present a menu. Run the searches yourself and answer from what they return. Empty results are a complete answer; an unanswered question is not.
|
||||
|
||||
## NEVER FABRICATE
|
||||
|
||||
|
|
|
|||
|
|
@ -20,7 +20,7 @@ from collections.abc import Callable
|
|||
from typing import Any
|
||||
|
||||
from src import crud
|
||||
from src.config import ReasoningLevel, settings
|
||||
from src.config import DialecticLevelSettings, ReasoningLevel, settings
|
||||
from src.dependencies import tracked_db
|
||||
from src.dialectic import prompts
|
||||
from src.dialectic.core import DialecticAgent
|
||||
|
|
@ -161,6 +161,30 @@ class WorkspaceDialecticAgent(DialecticAgent):
|
|||
tools = [t for t in tools if t.get("name") not in unscopable]
|
||||
return tools
|
||||
|
||||
def _tool_choice(
|
||||
self, level_settings: DialecticLevelSettings
|
||||
) -> str | dict[str, Any] | None:
|
||||
"""Require a tool call on the first turn.
|
||||
|
||||
The pair agent prefetches the observations relevant to its query, so it
|
||||
can legitimately answer from context alone. This agent's prefetch is an
|
||||
orientation overview — scale, active peers, their cards — not the corpus.
|
||||
Left free to skip tools, the model treats that overview as everything it
|
||||
has: it answers when the overview happens to carry the fact, and
|
||||
otherwise writes out the search it should have run and asks the caller
|
||||
which option to take. Workspace chat has no caller to answer, so that
|
||||
response is dead on arrival.
|
||||
|
||||
Recall is the job, so make the first search mandatory and let the loop
|
||||
relax to "auto" afterwards. Any other value a level configures is passed
|
||||
through untouched, so this only overrides the two cases that let the
|
||||
model opt out entirely.
|
||||
"""
|
||||
choice = level_settings.TOOL_CHOICE
|
||||
if choice is None or choice == "auto":
|
||||
return "required"
|
||||
return choice
|
||||
|
||||
async def _create_tool_executor(self) -> Callable[[str, dict[str, Any]], Any]:
|
||||
return await create_workspace_tool_executor(
|
||||
workspace_name=self.workspace_name,
|
||||
|
|
|
|||
|
|
@ -0,0 +1,216 @@
|
|||
"""Read-only count of the collections whose next dream is due. Enqueues nothing."""
|
||||
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from logging import getLogger
|
||||
from typing import Any, cast
|
||||
|
||||
from sqlalchemy import func, select
|
||||
from sqlalchemy.dialects.postgresql import aggregate_order_by
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from src import models
|
||||
from src.config import settings
|
||||
from src.schemas import DreamType
|
||||
from src.utils.config_helpers import get_configuration
|
||||
from src.utils.work_unit import construct_work_unit_key
|
||||
|
||||
logger = getLogger(__name__)
|
||||
|
||||
|
||||
async def count_due_dreams(db: AsyncSession) -> int:
|
||||
"""Count collections past the threshold, the idle timeout, the min-hours gate, any earlier attempt, and the session's dream setting."""
|
||||
dream_types = [
|
||||
DreamType(dream_type)
|
||||
for dream_type in settings.DREAM.ENABLED_TYPES
|
||||
if dream_type == DreamType.OMNI.value
|
||||
]
|
||||
if not settings.DREAM.ENABLED or not dream_types:
|
||||
return 0
|
||||
|
||||
explicit_counts = (
|
||||
select(
|
||||
models.Document.workspace_name,
|
||||
models.Document.observer,
|
||||
models.Document.observed,
|
||||
func.count(models.Document.id).label("explicit_count"),
|
||||
func.max(models.Document.created_at).label("newest_created_at"),
|
||||
func.array_agg(
|
||||
aggregate_order_by(
|
||||
models.Document.session_name, models.Document.created_at.desc()
|
||||
)
|
||||
)[1].label("newest_session_name"),
|
||||
)
|
||||
.where(models.Document.level == "explicit")
|
||||
.group_by(
|
||||
models.Document.workspace_name,
|
||||
models.Document.observer,
|
||||
models.Document.observed,
|
||||
)
|
||||
.subquery()
|
||||
)
|
||||
|
||||
rows = (
|
||||
await db.execute(
|
||||
select(
|
||||
models.Collection.workspace_name,
|
||||
models.Collection.observer,
|
||||
models.Collection.observed,
|
||||
models.Collection.internal_metadata,
|
||||
func.coalesce(explicit_counts.c.explicit_count, 0),
|
||||
explicit_counts.c.newest_created_at,
|
||||
explicit_counts.c.newest_session_name,
|
||||
).outerjoin(
|
||||
explicit_counts,
|
||||
(models.Collection.workspace_name == explicit_counts.c.workspace_name)
|
||||
& (models.Collection.observer == explicit_counts.c.observer)
|
||||
& (models.Collection.observed == explicit_counts.c.observed),
|
||||
)
|
||||
)
|
||||
).all()
|
||||
|
||||
now = datetime.now(UTC)
|
||||
idle_cutoff = now - timedelta(minutes=settings.DREAM.IDLE_TIMEOUT_MINUTES)
|
||||
candidates: dict[str, tuple[str, str, datetime]] = {}
|
||||
|
||||
for row in rows:
|
||||
workspace_name = cast(str, row[0])
|
||||
observer = cast(str, row[1])
|
||||
observed = cast(str, row[2])
|
||||
internal_metadata = cast("dict[str, Any] | None", row[3])
|
||||
explicit_count = cast(int, row[4])
|
||||
newest_created_at = cast("datetime | None", row[5])
|
||||
newest_session_name = cast("str | None", row[6])
|
||||
|
||||
dream_metadata: dict[str, Any] = (internal_metadata or {}).get("dream", {})
|
||||
since_last_dream = explicit_count - int(
|
||||
dream_metadata.get("last_dream_document_count", 0)
|
||||
)
|
||||
if since_last_dream < settings.DREAM.DOCUMENT_THRESHOLD:
|
||||
continue
|
||||
|
||||
if newest_created_at is None or newest_created_at > idle_cutoff:
|
||||
continue
|
||||
|
||||
if newest_session_name is None:
|
||||
continue
|
||||
|
||||
last_dream_at = cast("str | None", dream_metadata.get("last_dream_at"))
|
||||
if last_dream_at and _within_min_hours_gate(last_dream_at, now):
|
||||
continue
|
||||
|
||||
for dream_type in dream_types:
|
||||
work_unit_key = construct_work_unit_key(
|
||||
workspace_name,
|
||||
{
|
||||
"task_type": "dream",
|
||||
"observer": observer,
|
||||
"observed": observed,
|
||||
"dream_type": dream_type.value,
|
||||
},
|
||||
)
|
||||
candidates[work_unit_key] = (
|
||||
workspace_name,
|
||||
newest_session_name,
|
||||
newest_created_at,
|
||||
)
|
||||
|
||||
if not candidates:
|
||||
return 0
|
||||
|
||||
attempt_rows = (
|
||||
await db.execute(
|
||||
select(
|
||||
models.QueueItem.work_unit_key,
|
||||
func.max(models.QueueItem.created_at),
|
||||
)
|
||||
.where(
|
||||
models.QueueItem.task_type == "dream",
|
||||
models.QueueItem.work_unit_key.in_(candidates.keys()),
|
||||
)
|
||||
.group_by(models.QueueItem.work_unit_key)
|
||||
)
|
||||
).all()
|
||||
newest_attempts: dict[str, datetime] = {
|
||||
cast(str, row[0]): cast(datetime, row[1]) for row in attempt_rows
|
||||
}
|
||||
|
||||
unattempted = [
|
||||
(workspace_name, session_name)
|
||||
for work_unit_key, (
|
||||
workspace_name,
|
||||
session_name,
|
||||
newest_created_at,
|
||||
) in candidates.items()
|
||||
if work_unit_key not in newest_attempts
|
||||
or newest_attempts[work_unit_key] < newest_created_at
|
||||
]
|
||||
if not unattempted:
|
||||
return 0
|
||||
|
||||
return await _count_with_dreams_enabled(db, unattempted)
|
||||
|
||||
|
||||
async def _count_with_dreams_enabled(
|
||||
db: AsyncSession, candidates: list[tuple[str, str]]
|
||||
) -> int:
|
||||
"""Drop candidates whose resolved configuration has dreams turned off."""
|
||||
workspace_names = {workspace_name for workspace_name, _ in candidates}
|
||||
session_keys = set(candidates)
|
||||
|
||||
workspaces = {
|
||||
workspace.name: workspace
|
||||
for workspace in (
|
||||
await db.execute(
|
||||
select(models.Workspace).where(
|
||||
models.Workspace.name.in_(workspace_names)
|
||||
)
|
||||
)
|
||||
)
|
||||
.scalars()
|
||||
.all()
|
||||
}
|
||||
|
||||
sessions: dict[tuple[str, str], models.Session] = {}
|
||||
if session_keys:
|
||||
session_rows = (
|
||||
(
|
||||
await db.execute(
|
||||
select(models.Session).where(
|
||||
models.Session.workspace_name.in_(workspace_names),
|
||||
models.Session.name.in_(
|
||||
{session_name for _, session_name in candidates}
|
||||
),
|
||||
)
|
||||
)
|
||||
)
|
||||
.scalars()
|
||||
.all()
|
||||
)
|
||||
sessions = {
|
||||
(session.workspace_name, session.name): session for session in session_rows
|
||||
}
|
||||
|
||||
enabled = 0
|
||||
for workspace_name, session_name in candidates:
|
||||
configuration = get_configuration(
|
||||
None,
|
||||
sessions.get((workspace_name, session_name)),
|
||||
workspaces.get(workspace_name),
|
||||
)
|
||||
if configuration.dream.enabled:
|
||||
enabled += 1
|
||||
return enabled
|
||||
|
||||
|
||||
def _within_min_hours_gate(last_dream_at: str, now: datetime) -> bool:
|
||||
"""True when the last dream is too recent for another one."""
|
||||
try:
|
||||
last_dream_time = datetime.fromisoformat(last_dream_at)
|
||||
except (ValueError, TypeError):
|
||||
return False
|
||||
|
||||
if last_dream_time.tzinfo is None:
|
||||
last_dream_time = last_dream_time.replace(tzinfo=UTC)
|
||||
|
||||
hours_since = (now - last_dream_time).total_seconds() / 3600
|
||||
return hours_since < settings.DREAM.MIN_HOURS_BETWEEN_DREAMS
|
||||
12
src/main.py
12
src/main.py
|
|
@ -15,6 +15,7 @@ from sentry_sdk.integrations.sqlalchemy import SqlalchemyIntegration
|
|||
from sentry_sdk.integrations.starlette import StarletteIntegration
|
||||
|
||||
from src._version import HONCHO_VERSION
|
||||
from src.backlog import DeriverMetricsPoller
|
||||
from src.cache.client import close_cache, init_cache
|
||||
from src.config import settings
|
||||
from src.db import (
|
||||
|
|
@ -26,6 +27,7 @@ from src.db import (
|
|||
from src.exceptions import HonchoException
|
||||
from src.routers import (
|
||||
conclusions,
|
||||
deriver_metrics,
|
||||
keys,
|
||||
messages,
|
||||
peers,
|
||||
|
|
@ -135,12 +137,21 @@ async def lifespan(_: FastAPI):
|
|||
"Error initializing cache in api process; proceeding without cache: %s", e
|
||||
)
|
||||
|
||||
deriver_metrics_poller = DeriverMetricsPoller()
|
||||
deriver_metrics.set_deriver_metrics_poller(deriver_metrics_poller)
|
||||
try:
|
||||
await deriver_metrics_poller.start()
|
||||
except Exception as e:
|
||||
logger.error("Failed to start backlog metrics poller: %s", e)
|
||||
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
# Import here to avoid circular import at module load time
|
||||
from src.vector_store import close_external_vector_store
|
||||
|
||||
await deriver_metrics_poller.shutdown()
|
||||
deriver_metrics.set_deriver_metrics_poller(None)
|
||||
await close_external_vector_store()
|
||||
await close_cache()
|
||||
await engine.dispose()
|
||||
|
|
@ -189,6 +200,7 @@ app.include_router(messages.router, prefix="/v3")
|
|||
app.include_router(conclusions.router, prefix="/v3")
|
||||
app.include_router(keys.router, prefix="/v3")
|
||||
app.include_router(webhooks.router, prefix="/v3")
|
||||
app.include_router(deriver_metrics.router)
|
||||
|
||||
# Prometheus metrics endpoint
|
||||
app.add_route("/metrics", metrics_endpoint, methods=["GET"])
|
||||
|
|
|
|||
|
|
@ -33,7 +33,7 @@ from src.dependencies import tracked_db
|
|||
from src.embedding_client import embedding_client
|
||||
from src.exceptions import VectorStoreError
|
||||
from src.reconciler.sync_vectors import (
|
||||
_backoff_eligible, # pyright: ignore[reportPrivateUsage]
|
||||
backoff_eligible,
|
||||
build_message_vector_record,
|
||||
compute_chunk_positions,
|
||||
)
|
||||
|
|
@ -177,7 +177,7 @@ async def _claim_and_lease(message_ids: list[str]) -> list[_ClaimedChunk]:
|
|||
and_(
|
||||
models.MessageEmbedding.message_id.in_(message_ids),
|
||||
models.MessageEmbedding.sync_state == "pending",
|
||||
_backoff_eligible(models.MessageEmbedding.last_sync_at),
|
||||
backoff_eligible(models.MessageEmbedding.last_sync_at),
|
||||
)
|
||||
)
|
||||
.order_by(models.MessageEmbedding.message_id, models.MessageEmbedding.id)
|
||||
|
|
|
|||
|
|
@ -39,7 +39,7 @@ MAX_SYNC_ATTEMPTS = 20 # After this many failures, mark as failed
|
|||
SYNC_BACKOFF = datetime.timedelta(minutes=10)
|
||||
|
||||
|
||||
def _backoff_eligible(
|
||||
def backoff_eligible(
|
||||
last_sync_at: InstrumentedAttribute[datetime.datetime | None],
|
||||
) -> ColumnElement[bool]:
|
||||
"""Rows are eligible for sync if never attempted or past the backoff window."""
|
||||
|
|
@ -92,7 +92,7 @@ async def _get_documents_needing_sync(
|
|||
and_(
|
||||
models.Document.deleted_at.is_(None),
|
||||
models.Document.sync_state == "pending", # Only pending items
|
||||
_backoff_eligible(models.Document.last_sync_at),
|
||||
backoff_eligible(models.Document.last_sync_at),
|
||||
)
|
||||
)
|
||||
.order_by(models.Document.last_sync_at.asc().nullsfirst())
|
||||
|
|
@ -132,7 +132,7 @@ async def _get_message_embeddings_needing_sync(
|
|||
.where(
|
||||
and_(
|
||||
models.MessageEmbedding.sync_state == "pending",
|
||||
_backoff_eligible(models.MessageEmbedding.last_sync_at),
|
||||
backoff_eligible(models.MessageEmbedding.last_sync_at),
|
||||
)
|
||||
)
|
||||
.group_by(models.MessageEmbedding.message_id)
|
||||
|
|
@ -153,7 +153,7 @@ async def _get_message_embeddings_needing_sync(
|
|||
and_(
|
||||
models.MessageEmbedding.message_id.in_(message_ids),
|
||||
models.MessageEmbedding.sync_state == "pending",
|
||||
_backoff_eligible(models.MessageEmbedding.last_sync_at),
|
||||
backoff_eligible(models.MessageEmbedding.last_sync_at),
|
||||
)
|
||||
)
|
||||
.order_by(models.MessageEmbedding.message_id, models.MessageEmbedding.id)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,41 @@
|
|||
"""Deriver work metrics as JSON, with the age of the measurement alongside them."""
|
||||
|
||||
from logging import getLogger
|
||||
|
||||
from fastapi import APIRouter, HTTPException
|
||||
|
||||
from src.backlog import DeriverMetricsPoller
|
||||
|
||||
logger = getLogger(__name__)
|
||||
|
||||
router = APIRouter(prefix="/deriver", tags=["deriver"])
|
||||
|
||||
_poller: DeriverMetricsPoller | None = None
|
||||
|
||||
|
||||
def set_deriver_metrics_poller(poller: DeriverMetricsPoller | None) -> None:
|
||||
global _poller
|
||||
_poller = poller
|
||||
|
||||
|
||||
@router.get("/metrics")
|
||||
async def get_deriver_metrics_response() -> dict[str, float | int]:
|
||||
"""Seconds of outstanding deriver work, plus the raw counts behind it."""
|
||||
snapshot = _poller.snapshot if _poller is not None else None
|
||||
if snapshot is None or snapshot.measured_at is None:
|
||||
raise HTTPException(
|
||||
status_code=503, detail="No deriver measurement available yet"
|
||||
)
|
||||
|
||||
return {
|
||||
"outstanding_work_seconds": snapshot.signal_seconds,
|
||||
"eligible_work_units": snapshot.stats.eligible_work_units,
|
||||
"claimed_work_units": snapshot.stats.claimed_work_units,
|
||||
"pending_items": snapshot.stats.pending_items,
|
||||
"oldest_pending_age_seconds": snapshot.stats.oldest_pending_age_seconds,
|
||||
"embeddings_pending": snapshot.stats.embeddings_pending,
|
||||
"embeddings_pending_due": snapshot.stats.embeddings_pending_due,
|
||||
"dreams_due": snapshot.dreams_due,
|
||||
"measured_at": snapshot.measured_at,
|
||||
"measurement_age_seconds": snapshot.age_seconds or 0.0,
|
||||
}
|
||||
|
|
@ -78,6 +78,7 @@ from src.schemas.configuration import (
|
|||
WorkspaceConfiguration,
|
||||
)
|
||||
from src.schemas.internal import (
|
||||
DeriverMetrics,
|
||||
DocumentBase,
|
||||
DocumentCreate,
|
||||
DocumentMetadata,
|
||||
|
|
@ -163,6 +164,7 @@ __all__ = [
|
|||
"WorkspaceMessageSearchOptions",
|
||||
"WorkspaceUpdate",
|
||||
# internal
|
||||
"DeriverMetrics",
|
||||
"DocumentBase",
|
||||
"DocumentCreate",
|
||||
"DocumentMetadata",
|
||||
|
|
|
|||
|
|
@ -31,6 +31,7 @@ from src.schemas.configuration import (
|
|||
SessionPeerConfig,
|
||||
WorkspaceConfiguration,
|
||||
)
|
||||
from src.utils.sanitization import NulStripped, strip_nul
|
||||
from src.utils.scopes import (
|
||||
SCOPE_PEER_PREFIX,
|
||||
is_scope_peer_name,
|
||||
|
|
@ -48,28 +49,6 @@ _METADATA_MAX_KEYS = 100
|
|||
_METADATA_MAX_DEPTH = 5
|
||||
|
||||
|
||||
def _sanitize_value(v: Any) -> Any:
|
||||
"""Recursively strip NUL bytes from strings in nested data structures."""
|
||||
if isinstance(v, str):
|
||||
return v.replace("\x00", "")
|
||||
if isinstance(v, dict):
|
||||
d = cast(dict[str, Any], v)
|
||||
return {_sanitize_value(k): _sanitize_value(val) for k, val in d.items()}
|
||||
if isinstance(v, list):
|
||||
lst = cast(list[Any], v)
|
||||
return [_sanitize_value(item) for item in lst]
|
||||
return v
|
||||
|
||||
|
||||
def _strip_nul(v: str) -> str:
|
||||
"""Strip NUL bytes from a string field (Postgres TEXT rejects \\x00)."""
|
||||
return v.replace("\x00", "")
|
||||
|
||||
|
||||
# Reusable annotation for query fields; composes with a per-field Field(...).
|
||||
NulStripped = AfterValidator(_strip_nul)
|
||||
|
||||
|
||||
def _check_metadata_limits(
|
||||
data: dict[str, Any],
|
||||
*,
|
||||
|
|
@ -97,7 +76,7 @@ def _validate_metadata(v: Any) -> Any:
|
|||
return v
|
||||
data = cast(dict[str, Any], v)
|
||||
_check_metadata_limits(data)
|
||||
return _sanitize_value(data)
|
||||
return strip_nul(data)
|
||||
|
||||
|
||||
_SanitizedMetadata = Annotated[dict[str, Any], BeforeValidator(_validate_metadata)]
|
||||
|
|
@ -331,7 +310,7 @@ class PeerCardSet(BaseModel):
|
|||
def sanitize_peer_card(cls, v: Any) -> Any:
|
||||
if isinstance(v, list):
|
||||
return [
|
||||
item.replace("\x00", "") if isinstance(item, str) else item
|
||||
strip_nul(item) if isinstance(item, str) else item
|
||||
for item in cast(list[Any], v)
|
||||
]
|
||||
return v
|
||||
|
|
@ -358,7 +337,7 @@ class MessageCreate(MessageBase):
|
|||
@field_validator("content", mode="after")
|
||||
@classmethod
|
||||
def sanitize_content(cls, v: str) -> str:
|
||||
return v.replace("\x00", "")
|
||||
return strip_nul(v)
|
||||
|
||||
@property
|
||||
def encoded_message(self) -> list[int]:
|
||||
|
|
@ -691,7 +670,7 @@ class ConclusionCreate(BaseModel):
|
|||
@field_validator("content", mode="after")
|
||||
@classmethod
|
||||
def sanitize_content(cls, v: str) -> str:
|
||||
return v.replace("\x00", "")
|
||||
return strip_nul(v)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_token_count(self) -> Self:
|
||||
|
|
|
|||
|
|
@ -6,10 +6,11 @@ These are not part of the public API contract and may change without notice.
|
|||
from enum import Enum
|
||||
from typing import Annotated, Literal, Self
|
||||
|
||||
from pydantic import BaseModel, Field, field_validator, model_validator
|
||||
from pydantic import BaseModel, Field, model_validator
|
||||
|
||||
from src.schemas.api import MessageCreate
|
||||
from src.schemas.configuration import SessionPeerConfig
|
||||
from src.utils.sanitization import NulStripped
|
||||
from src.utils.types import DocumentLevel
|
||||
|
||||
|
||||
|
|
@ -59,7 +60,7 @@ class DocumentMetadata(BaseModel):
|
|||
|
||||
|
||||
class DocumentCreate(DocumentBase):
|
||||
content: Annotated[str, Field(min_length=1, max_length=100000)]
|
||||
content: Annotated[str, Field(min_length=1, max_length=100000), NulStripped]
|
||||
session_name: str | None = Field(
|
||||
default=None,
|
||||
description="The session from which the document was derived (NULL for global observations)",
|
||||
|
|
@ -85,7 +86,7 @@ class DocumentCreate(DocumentBase):
|
|||
class ObservationInput(BaseModel):
|
||||
"""Validated observation input from LLM tool calls."""
|
||||
|
||||
content: Annotated[str, Field(min_length=1)]
|
||||
content: Annotated[str, Field(min_length=1), NulStripped]
|
||||
level: DocumentLevel = "explicit"
|
||||
source_ids: list[str] | None = None
|
||||
premises: list[str] | None = None
|
||||
|
|
@ -96,11 +97,6 @@ class ObservationInput(BaseModel):
|
|||
) = None
|
||||
confidence: Literal["high", "medium", "low"] | None = None
|
||||
|
||||
@field_validator("content", mode="after")
|
||||
@classmethod
|
||||
def sanitize_content(cls, v: str) -> str:
|
||||
return v.replace("\x00", "")
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_level_fields(self) -> Self:
|
||||
"""Validate that level-specific fields are present when required."""
|
||||
|
|
@ -144,6 +140,17 @@ class QueueCounts(BaseModel):
|
|||
sessions: dict[str, SessionCounts]
|
||||
|
||||
|
||||
class DeriverMetrics(BaseModel):
|
||||
"""Database-wide view of the deriver's outstanding work."""
|
||||
|
||||
eligible_work_units: int = 0
|
||||
claimed_work_units: int = 0
|
||||
pending_items: int = 0
|
||||
oldest_pending_age_seconds: float = 0.0
|
||||
embeddings_pending: int = 0
|
||||
embeddings_pending_due: int = 0
|
||||
|
||||
|
||||
class QueueStatusRow(BaseModel):
|
||||
"""Represents a row from the queue status SQL query result."""
|
||||
|
||||
|
|
|
|||
|
|
@ -199,6 +199,69 @@ message_embeddings_pending_gauge = NamespacedGauge(
|
|||
["namespace"],
|
||||
)
|
||||
|
||||
message_embeddings_pending_due_gauge = NamespacedGauge(
|
||||
"message_embeddings_pending_due",
|
||||
"Pending MessageEmbedding rows past their retry backoff, so a sync attempt "
|
||||
+ "is due. Service-wide DB count, reported independently by every API "
|
||||
+ "replica — aggregate with max() or avg(), never sum()",
|
||||
["namespace"],
|
||||
)
|
||||
|
||||
deriver_outstanding_work_seconds_gauge = NamespacedGauge(
|
||||
"deriver_outstanding_work_seconds",
|
||||
"Seconds of outstanding deriver work, 0 when a deriver has nothing to do. "
|
||||
+ "Service-wide DB value, reported independently by every API replica — "
|
||||
+ "aggregate with max(), never sum()",
|
||||
["namespace"],
|
||||
)
|
||||
|
||||
deriver_queue_work_units_eligible_gauge = NamespacedGauge(
|
||||
"deriver_queue_work_units_eligible",
|
||||
"Work units a deriver could claim right now, ignoring stale claims. "
|
||||
+ "Service-wide DB count, reported independently by every API replica — "
|
||||
+ "aggregate with max() or avg(), never sum()",
|
||||
["namespace"],
|
||||
)
|
||||
|
||||
deriver_queue_work_units_claimed_gauge = NamespacedGauge(
|
||||
"deriver_queue_work_units_claimed",
|
||||
"Work units held by a claim refreshed inside the stale timeout, so work is "
|
||||
+ "in flight. Service-wide DB count, reported independently by every API "
|
||||
+ "replica — aggregate with max() or avg(), never sum()",
|
||||
["namespace"],
|
||||
)
|
||||
|
||||
deriver_queue_items_pending_gauge = NamespacedGauge(
|
||||
"deriver_queue_items_pending",
|
||||
"Unprocessed queue rows, whether or not they are claimable yet. "
|
||||
+ "Service-wide DB count, reported independently by every API replica — "
|
||||
+ "aggregate with max() or avg(), never sum()",
|
||||
["namespace"],
|
||||
)
|
||||
|
||||
deriver_queue_oldest_pending_age_seconds_gauge = NamespacedGauge(
|
||||
"deriver_queue_oldest_pending_age_seconds",
|
||||
"Age of the oldest unprocessed queue row, 0 when the queue is empty. "
|
||||
+ "Service-wide DB value, reported independently by every API replica — "
|
||||
+ "aggregate with max() or avg(), never sum()",
|
||||
["namespace"],
|
||||
)
|
||||
|
||||
dreams_due_gauge = NamespacedGauge(
|
||||
"dreams_due",
|
||||
"Collections whose next dream is due and would actually run. "
|
||||
+ "Service-wide DB count, reported independently by every API replica — "
|
||||
+ "aggregate with max() or avg(), never sum()",
|
||||
["namespace"],
|
||||
)
|
||||
|
||||
deriver_metrics_last_success_timestamp_gauge = NamespacedGauge(
|
||||
"deriver_metrics_last_success_timestamp_seconds",
|
||||
"Unix time of the last successful deriver-metrics refresh in this replica. "
|
||||
+ "Alert on time() minus this value; a frozen value means the poller stopped",
|
||||
["namespace"],
|
||||
)
|
||||
|
||||
# DB connection-pool health. The in-flight gauge counts statements actually
|
||||
# executing on the wire, so checked_out minus in_flight reveals connections held
|
||||
# but parked (the "idle in transaction during an external call" antipattern).
|
||||
|
|
@ -508,6 +571,10 @@ class PrometheusMetrics:
|
|||
self._touch(embed_now_tasks_shed_counter)
|
||||
self.set_embed_now_tasks_in_flight(0)
|
||||
|
||||
self.set_deriver_metrics()
|
||||
self.set_deriver_outstanding_work(seconds=0)
|
||||
self.set_dreams_due(count=0)
|
||||
|
||||
elif instance_type == "deriver":
|
||||
# deriver tokens: only the valid (token_type, component) tuples per
|
||||
# task_type (see _DERIVER_TOKEN_COMBOS_BY_TASK).
|
||||
|
|
@ -548,6 +615,46 @@ class PrometheusMetrics:
|
|||
except Exception as e:
|
||||
self._handle_metric_error("set_message_embeddings_pending", e)
|
||||
|
||||
def set_deriver_metrics(
|
||||
self,
|
||||
*,
|
||||
eligible_work_units: int = 0,
|
||||
claimed_work_units: int = 0,
|
||||
pending_items: int = 0,
|
||||
oldest_pending_age_seconds: float = 0.0,
|
||||
embeddings_pending: int = 0,
|
||||
embeddings_pending_due: int = 0,
|
||||
) -> None:
|
||||
try:
|
||||
deriver_queue_work_units_eligible_gauge.labels().set(eligible_work_units)
|
||||
deriver_queue_work_units_claimed_gauge.labels().set(claimed_work_units)
|
||||
deriver_queue_items_pending_gauge.labels().set(pending_items)
|
||||
deriver_queue_oldest_pending_age_seconds_gauge.labels().set(
|
||||
oldest_pending_age_seconds
|
||||
)
|
||||
message_embeddings_pending_gauge.labels().set(embeddings_pending)
|
||||
message_embeddings_pending_due_gauge.labels().set(embeddings_pending_due)
|
||||
except Exception as e:
|
||||
self._handle_metric_error("set_deriver_metrics", e)
|
||||
|
||||
def set_deriver_outstanding_work(self, *, seconds: float) -> None:
|
||||
try:
|
||||
deriver_outstanding_work_seconds_gauge.labels().set(seconds)
|
||||
except Exception as e:
|
||||
self._handle_metric_error("set_deriver_outstanding_work", e)
|
||||
|
||||
def set_dreams_due(self, *, count: int) -> None:
|
||||
try:
|
||||
dreams_due_gauge.labels().set(count)
|
||||
except Exception as e:
|
||||
self._handle_metric_error("set_dreams_due", e)
|
||||
|
||||
def set_deriver_metrics_last_success(self, *, timestamp: float) -> None:
|
||||
try:
|
||||
deriver_metrics_last_success_timestamp_gauge.labels().set(timestamp)
|
||||
except Exception as e:
|
||||
self._handle_metric_error("set_deriver_metrics_last_success", e)
|
||||
|
||||
|
||||
prometheus_metrics = PrometheusMetrics()
|
||||
|
||||
|
|
|
|||
|
|
@ -36,6 +36,7 @@ from src.utils.representation import (
|
|||
Representation,
|
||||
allowlist_safe_levels,
|
||||
)
|
||||
from src.utils.sanitization import strip_nul
|
||||
from src.utils.types import ToolResult, embedding_call_purpose, get_current_iteration
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
|
@ -77,8 +78,20 @@ def _validate_peer_card_entry(line: str) -> bool:
|
|||
def _normalized_observation_input(
|
||||
obs: schemas.ObservationInput,
|
||||
) -> schemas.ObservationInput:
|
||||
"""Return an observation input with content normalized for persistence/embedding."""
|
||||
return obs.model_copy(update={"content": obs.content.strip()})
|
||||
"""Return an observation input with content normalized for persistence/embedding.
|
||||
|
||||
NUL bytes are removed here rather than closer to the database so that the
|
||||
text that gets embedded is the same text that gets stored. `premises` and
|
||||
`sources` ride along in internal_metadata, and jsonb rejects NUL in strings
|
||||
just as text columns do.
|
||||
"""
|
||||
return obs.model_copy(
|
||||
update={
|
||||
"content": strip_nul(obs.content).strip(),
|
||||
"premises": strip_nul(obs.premises),
|
||||
"sources": strip_nul(obs.sources),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _base_observation_properties() -> dict[str, Any]:
|
||||
|
|
@ -986,10 +999,12 @@ async def create_observations(
|
|||
logger.warning("create_observations called with empty list")
|
||||
return ObservationsCreatedResult(created_count=0, created_levels=[], failed=[])
|
||||
|
||||
# Normalize before the emptiness check: str.strip() does not remove NUL,
|
||||
# so content that normalizes away has to be dropped afterwards.
|
||||
normalized_observations = [
|
||||
_normalized_observation_input(obs)
|
||||
for obs in observations
|
||||
if obs.content.strip()
|
||||
normalized
|
||||
for normalized in (_normalized_observation_input(obs) for obs in observations)
|
||||
if normalized.content
|
||||
]
|
||||
if not normalized_observations:
|
||||
logger.info("No non-empty observations to create")
|
||||
|
|
|
|||
|
|
@ -5,6 +5,14 @@ from pydantic import BaseModel, ConfigDict
|
|||
|
||||
from src.schemas import DreamType, ReconcilerType, ResolvedConfiguration
|
||||
|
||||
# Queue mechanics, not task data: the deriver stores a per-work-unit transient
|
||||
# failure count under this key so a retry budget survives work-unit reclaim.
|
||||
# Every payload model below forbids extras, so anything that reads a raw
|
||||
# QueueItem.payload must strip this key before validating. Lives here rather
|
||||
# than in the deriver because both the writer (queue_manager) and the stripper
|
||||
# (consumer) need it, and queue_manager imports consumer.
|
||||
RETRY_ATTEMPTS_PAYLOAD_KEY = "_retry_attempts"
|
||||
|
||||
|
||||
class BasePayload(BaseModel):
|
||||
"""Base payload with common fields."""
|
||||
|
|
|
|||
|
|
@ -0,0 +1,86 @@
|
|||
"""Classify exceptions as transient (safe to retry) or terminal.
|
||||
|
||||
Imports only exception taxonomies, so it is importable from anywhere and
|
||||
unit-testable without a DB.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Iterator
|
||||
|
||||
import httpx
|
||||
from sqlalchemy.exc import DBAPIError
|
||||
|
||||
__all__ = ["is_retryable_db_error", "is_retryable_error"]
|
||||
|
||||
_RETRYABLE_SQLSTATES = frozenset(
|
||||
{
|
||||
"40001", # serialization_failure
|
||||
"40P01", # deadlock_detected
|
||||
"55P03", # lock_not_available (lock_timeout / NOWAIT)
|
||||
"57014", # query_canceled (statement_timeout)
|
||||
"08000", # connection_exception family
|
||||
"08001",
|
||||
"08003",
|
||||
"08004",
|
||||
"08006",
|
||||
}
|
||||
)
|
||||
|
||||
# Provider/network transport failures. SDK wrappers (anthropic/openai
|
||||
# APIConnectionError etc.) chain to these via __cause__.
|
||||
_TRANSPORT_ERRORS = (
|
||||
httpx.TransportError,
|
||||
ConnectionError,
|
||||
asyncio.TimeoutError,
|
||||
TimeoutError,
|
||||
)
|
||||
|
||||
|
||||
def _iter_cause_chain(exc: BaseException) -> Iterator[BaseException]:
|
||||
seen: set[int] = set()
|
||||
current: BaseException | None = exc
|
||||
while current is not None and id(current) not in seen:
|
||||
seen.add(id(current))
|
||||
yield current
|
||||
current = current.__cause__
|
||||
|
||||
|
||||
def _sqlstate(exc: DBAPIError) -> str | None:
|
||||
"""Extract the SQLSTATE off ``DBAPIError.orig``, driver-agnostically."""
|
||||
orig = getattr(exc, "orig", None)
|
||||
for candidate in (orig, getattr(orig, "__cause__", None)):
|
||||
code = getattr(candidate, "sqlstate", None)
|
||||
if isinstance(code, str):
|
||||
return code
|
||||
return None
|
||||
|
||||
|
||||
def is_retryable_db_error(exc: BaseException) -> bool:
|
||||
"""True for transient DB failures: deadlock, serialization failure,
|
||||
lock/statement timeout, or a lost connection.
|
||||
|
||||
Integrity (23xxx), data (22xxx), and programming (42xxx) errors are
|
||||
deliberately terminal.
|
||||
"""
|
||||
for current in _iter_cause_chain(exc):
|
||||
if not isinstance(current, DBAPIError):
|
||||
continue
|
||||
if current.connection_invalidated:
|
||||
return True
|
||||
if _sqlstate(current) in _RETRYABLE_SQLSTATES:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def is_retryable_error(exc: BaseException) -> bool:
|
||||
"""Superset of ``is_retryable_db_error``: also transient network/provider
|
||||
transport failures (timeouts, connection refused/reset).
|
||||
|
||||
Auth failures (401 from a rotated key) are deliberately terminal: they
|
||||
never self-heal, so retrying only delays the burn.
|
||||
"""
|
||||
if is_retryable_db_error(exc):
|
||||
return True
|
||||
return any(
|
||||
isinstance(current, _TRANSPORT_ERRORS) for current in _iter_cause_chain(exc)
|
||||
)
|
||||
|
|
@ -0,0 +1,46 @@
|
|||
"""Helpers for stripping bytes Postgres cannot store in text columns.
|
||||
|
||||
Postgres rejects NUL (0x00) in ``text``/``varchar`` values and in ``jsonb``
|
||||
strings, so any string bound into a query or persisted to those columns has to
|
||||
have NUL removed first. This applies to model-generated text as much as to
|
||||
user-supplied input: an LLM can emit a ``\\u0000`` escape in its tool-call
|
||||
arguments, which the JSON parser decodes into a real NUL byte.
|
||||
"""
|
||||
|
||||
from typing import Any, cast, overload
|
||||
|
||||
from pydantic import BeforeValidator
|
||||
|
||||
__all__ = ["NulStripped", "strip_nul"]
|
||||
|
||||
|
||||
@overload
|
||||
def strip_nul(value: str) -> str: ...
|
||||
|
||||
|
||||
@overload
|
||||
def strip_nul(value: Any) -> Any: ...
|
||||
|
||||
|
||||
def strip_nul(value: Any) -> Any:
|
||||
"""Recursively remove NUL bytes from strings, including nested ones.
|
||||
|
||||
Dict keys are stripped alongside values. Anything that is not a string,
|
||||
dict, or list -- ``None`` included -- is returned unchanged, so this can be
|
||||
applied to an optional field without a guard.
|
||||
"""
|
||||
if isinstance(value, str):
|
||||
return value.replace("\x00", "")
|
||||
if isinstance(value, dict):
|
||||
d = cast(dict[str, Any], value)
|
||||
return {strip_nul(k): strip_nul(v) for k, v in d.items()}
|
||||
if isinstance(value, list):
|
||||
lst = cast(list[Any], value)
|
||||
return [strip_nul(item) for item in lst]
|
||||
return value
|
||||
|
||||
|
||||
# Reusable annotation for string fields; composes with a per-field Field(...).
|
||||
# Runs *before* the field's own constraints, so `min_length` is checked against
|
||||
# the stripped value and all-NUL input is rejected instead of becoming "".
|
||||
NulStripped = BeforeValidator(strip_nul)
|
||||
|
|
@ -0,0 +1,394 @@
|
|||
import datetime
|
||||
|
||||
import pytest
|
||||
from nanoid import generate as generate_nanoid
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from src import crud, models
|
||||
from src.config import settings
|
||||
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
|
||||
async def _make_session(
|
||||
db: AsyncSession, workspace: models.Workspace
|
||||
) -> models.Session:
|
||||
session = models.Session(name=str(generate_nanoid()), workspace_name=workspace.name)
|
||||
db.add(session)
|
||||
await db.flush()
|
||||
return session
|
||||
|
||||
|
||||
async def _add_representation_item(
|
||||
db: AsyncSession,
|
||||
workspace: models.Workspace,
|
||||
peer: models.Peer,
|
||||
session: models.Session,
|
||||
*,
|
||||
work_unit_key: str,
|
||||
token_count: int,
|
||||
age_seconds: int = 0,
|
||||
seq: int = 1,
|
||||
) -> models.QueueItem:
|
||||
message = models.Message(
|
||||
session_name=session.name,
|
||||
content="x",
|
||||
token_count=token_count,
|
||||
seq_in_session=seq,
|
||||
peer_name=peer.name,
|
||||
workspace_name=workspace.name,
|
||||
)
|
||||
db.add(message)
|
||||
await db.flush()
|
||||
|
||||
item = models.QueueItem(
|
||||
session_id=session.id,
|
||||
work_unit_key=work_unit_key,
|
||||
task_type="representation",
|
||||
payload={},
|
||||
processed=False,
|
||||
workspace_name=workspace.name,
|
||||
message_id=message.id,
|
||||
created_at=datetime.datetime.now(datetime.UTC)
|
||||
- datetime.timedelta(seconds=age_seconds),
|
||||
)
|
||||
db.add(item)
|
||||
await db.flush()
|
||||
return item
|
||||
|
||||
|
||||
async def _add_message(
|
||||
db: AsyncSession,
|
||||
workspace: models.Workspace,
|
||||
peer: models.Peer,
|
||||
session: models.Session,
|
||||
*,
|
||||
seq: int = 1,
|
||||
) -> models.Message:
|
||||
message = models.Message(
|
||||
session_name=session.name,
|
||||
content="x",
|
||||
token_count=1,
|
||||
seq_in_session=seq,
|
||||
peer_name=peer.name,
|
||||
workspace_name=workspace.name,
|
||||
)
|
||||
db.add(message)
|
||||
await db.flush()
|
||||
return message
|
||||
|
||||
|
||||
def _stale_timestamp() -> datetime.datetime:
|
||||
return datetime.datetime.now(datetime.UTC) - datetime.timedelta(
|
||||
minutes=settings.DERIVER.STALE_SESSION_TIMEOUT_MINUTES + 1
|
||||
)
|
||||
|
||||
|
||||
class TestDeriverMetrics:
|
||||
async def test_empty_queue_reports_zero(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer], # pyright: ignore[reportUnusedParameter]
|
||||
):
|
||||
stats = await crud.get_deriver_metrics(db_session)
|
||||
|
||||
assert stats.eligible_work_units == 0
|
||||
assert stats.claimed_work_units == 0
|
||||
assert stats.pending_items == 0
|
||||
assert stats.oldest_pending_age_seconds == 0.0
|
||||
|
||||
async def test_sub_threshold_batch_is_pending_but_not_eligible(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
"""A small, fresh batch is real work that a deriver would not yet claim."""
|
||||
workspace, peer = sample_data
|
||||
session = await _make_session(db_session, workspace)
|
||||
|
||||
await _add_representation_item(
|
||||
db_session,
|
||||
workspace,
|
||||
peer,
|
||||
session,
|
||||
work_unit_key="representation:small",
|
||||
token_count=1,
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
stats = await crud.get_deriver_metrics(db_session)
|
||||
|
||||
assert stats.pending_items == 1
|
||||
assert stats.eligible_work_units == 0
|
||||
|
||||
async def test_token_threshold_makes_batch_eligible(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
workspace, peer = sample_data
|
||||
session = await _make_session(db_session, workspace)
|
||||
|
||||
await _add_representation_item(
|
||||
db_session,
|
||||
workspace,
|
||||
peer,
|
||||
session,
|
||||
work_unit_key="representation:big",
|
||||
token_count=settings.DERIVER.REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS,
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
stats = await crud.get_deriver_metrics(db_session)
|
||||
|
||||
assert stats.eligible_work_units == 1
|
||||
|
||||
async def test_age_flush_makes_sub_threshold_batch_eligible(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
workspace, peer = sample_data
|
||||
session = await _make_session(db_session, workspace)
|
||||
|
||||
await _add_representation_item(
|
||||
db_session,
|
||||
workspace,
|
||||
peer,
|
||||
session,
|
||||
work_unit_key="representation:old",
|
||||
token_count=1,
|
||||
age_seconds=settings.DERIVER.REPRESENTATION_BATCH_MAX_AGE_SECONDS + 60,
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
stats = await crud.get_deriver_metrics(db_session)
|
||||
|
||||
assert stats.eligible_work_units == 1
|
||||
assert stats.oldest_pending_age_seconds >= (
|
||||
settings.DERIVER.REPRESENTATION_BATCH_MAX_AGE_SECONDS
|
||||
)
|
||||
|
||||
async def test_non_representation_work_is_eligible_immediately(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
workspace, _peer = sample_data
|
||||
|
||||
db_session.add(
|
||||
models.QueueItem(
|
||||
work_unit_key="reconciler:sync_vectors",
|
||||
task_type="reconciler",
|
||||
payload={},
|
||||
processed=False,
|
||||
workspace_name=workspace.name,
|
||||
)
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
stats = await crud.get_deriver_metrics(db_session)
|
||||
|
||||
assert stats.eligible_work_units == 1
|
||||
|
||||
async def test_live_claim_is_counted_as_work_in_flight(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
"""A claimed work unit is not claimable, but it is still outstanding work."""
|
||||
workspace, peer = sample_data
|
||||
session = await _make_session(db_session, workspace)
|
||||
|
||||
await _add_representation_item(
|
||||
db_session,
|
||||
workspace,
|
||||
peer,
|
||||
session,
|
||||
work_unit_key="representation:claimed",
|
||||
token_count=settings.DERIVER.REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS,
|
||||
)
|
||||
db_session.add(
|
||||
models.ActiveQueueSession(work_unit_key="representation:claimed")
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
stats = await crud.get_deriver_metrics(db_session)
|
||||
|
||||
assert stats.eligible_work_units == 0
|
||||
assert stats.claimed_work_units == 1
|
||||
|
||||
async def test_stale_claim_does_not_hide_work_and_is_not_in_flight(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
"""A dead worker's claim must not read as in flight, and must not hide work."""
|
||||
workspace, peer = sample_data
|
||||
session = await _make_session(db_session, workspace)
|
||||
|
||||
await _add_representation_item(
|
||||
db_session,
|
||||
workspace,
|
||||
peer,
|
||||
session,
|
||||
work_unit_key="representation:abandoned",
|
||||
token_count=settings.DERIVER.REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS,
|
||||
)
|
||||
db_session.add(
|
||||
models.ActiveQueueSession(
|
||||
work_unit_key="representation:abandoned",
|
||||
last_updated=_stale_timestamp(),
|
||||
)
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
stats = await crud.get_deriver_metrics(db_session)
|
||||
|
||||
assert stats.eligible_work_units == 1
|
||||
assert stats.claimed_work_units == 0
|
||||
|
||||
async def test_processed_items_are_not_counted(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
workspace, peer = sample_data
|
||||
session = await _make_session(db_session, workspace)
|
||||
|
||||
item = await _add_representation_item(
|
||||
db_session,
|
||||
workspace,
|
||||
peer,
|
||||
session,
|
||||
work_unit_key="representation:done",
|
||||
token_count=settings.DERIVER.REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS,
|
||||
)
|
||||
item.processed = True
|
||||
await db_session.commit()
|
||||
|
||||
stats = await crud.get_deriver_metrics(db_session)
|
||||
|
||||
assert stats.pending_items == 0
|
||||
assert stats.eligible_work_units == 0
|
||||
assert stats.oldest_pending_age_seconds == 0.0
|
||||
|
||||
|
||||
class TestPendingEmbeddings:
|
||||
async def test_never_attempted_row_is_due(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
workspace, peer = sample_data
|
||||
session = await _make_session(db_session, workspace)
|
||||
message = await _add_message(db_session, workspace, peer, session)
|
||||
db_session.add(
|
||||
models.MessageEmbedding(
|
||||
content="x",
|
||||
message_id=message.public_id,
|
||||
workspace_name=workspace.name,
|
||||
session_name=session.name,
|
||||
peer_name=peer.name,
|
||||
sync_state="pending",
|
||||
)
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
stats = await crud.get_deriver_metrics(db_session)
|
||||
|
||||
assert stats.embeddings_pending == 1
|
||||
assert stats.embeddings_pending_due == 1
|
||||
|
||||
async def test_row_inside_its_retry_wait_is_pending_but_not_due(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
"""A backing-off row is work the deriver cannot act on yet."""
|
||||
workspace, peer = sample_data
|
||||
session = await _make_session(db_session, workspace)
|
||||
message = await _add_message(db_session, workspace, peer, session)
|
||||
db_session.add(
|
||||
models.MessageEmbedding(
|
||||
content="x",
|
||||
message_id=message.public_id,
|
||||
workspace_name=workspace.name,
|
||||
session_name=session.name,
|
||||
peer_name=peer.name,
|
||||
sync_state="pending",
|
||||
last_sync_at=datetime.datetime.now(datetime.UTC),
|
||||
sync_attempts=1,
|
||||
)
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
stats = await crud.get_deriver_metrics(db_session)
|
||||
|
||||
assert stats.embeddings_pending == 1
|
||||
assert stats.embeddings_pending_due == 0
|
||||
|
||||
async def test_synced_rows_are_not_counted(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
workspace, peer = sample_data
|
||||
session = await _make_session(db_session, workspace)
|
||||
message = await _add_message(db_session, workspace, peer, session)
|
||||
db_session.add(
|
||||
models.MessageEmbedding(
|
||||
content="x",
|
||||
message_id=message.public_id,
|
||||
workspace_name=workspace.name,
|
||||
session_name=session.name,
|
||||
peer_name=peer.name,
|
||||
sync_state="synced",
|
||||
)
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
stats = await crud.get_deriver_metrics(db_session)
|
||||
|
||||
assert stats.embeddings_pending == 0
|
||||
assert stats.embeddings_pending_due == 0
|
||||
|
||||
|
||||
class TestMetricsAgreeWithDeriver:
|
||||
@pytest.mark.parametrize(
|
||||
"token_count,age_seconds",
|
||||
[
|
||||
(1, 0),
|
||||
(settings.DERIVER.REPRESENTATION_BATCH_WORK_UNIT_TARGET_TOKENS, 0),
|
||||
(1, settings.DERIVER.REPRESENTATION_BATCH_MAX_AGE_SECONDS + 60),
|
||||
],
|
||||
ids=["sub-threshold", "token-threshold", "age-flush"],
|
||||
)
|
||||
async def test_eligible_count_matches_what_the_deriver_claims(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
token_count: int,
|
||||
age_seconds: int,
|
||||
):
|
||||
"""The gauge is only trustworthy if it uses the deriver's own rule."""
|
||||
from src.deriver.queue_manager import QueueManager
|
||||
|
||||
workspace, peer = sample_data
|
||||
session = await _make_session(db_session, workspace)
|
||||
|
||||
await _add_representation_item(
|
||||
db_session,
|
||||
workspace,
|
||||
peer,
|
||||
session,
|
||||
work_unit_key="representation:agreement",
|
||||
token_count=token_count,
|
||||
age_seconds=age_seconds,
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
expected = (await crud.get_deriver_metrics(db_session)).eligible_work_units
|
||||
claimed = await QueueManager().get_and_claim_work_units()
|
||||
|
||||
assert len(claimed) == expected
|
||||
|
|
@ -1,10 +1,13 @@
|
|||
import asyncio
|
||||
import datetime
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from nanoid import generate as generate_nanoid
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.exc import OperationalError
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker
|
||||
|
||||
from src import crud, models, schemas
|
||||
from src.crud.document import SemanticRejectionResult, is_rejected_duplicate
|
||||
|
|
@ -195,7 +198,7 @@ class TestDocumentCRUD:
|
|||
deleted_doc = docs["User likes pizza"]
|
||||
kept_doc = docs["User dislikes vegetables"]
|
||||
|
||||
deleted_doc.deleted_at = datetime.datetime.now(datetime.timezone.utc)
|
||||
deleted_doc.deleted_at = datetime.datetime.now(datetime.UTC)
|
||||
await db_session.commit()
|
||||
|
||||
results = await crud.query_documents(
|
||||
|
|
@ -290,7 +293,7 @@ class TestDocumentCRUD:
|
|||
db_session, test_workspace, test_peer
|
||||
)
|
||||
|
||||
base = datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc)
|
||||
base = datetime.datetime(2026, 1, 1, tzinfo=datetime.UTC)
|
||||
# Three conclusions, all reinforced once -- the real-world steady state
|
||||
# before the fix -- inserted oldest-first.
|
||||
for i in range(3):
|
||||
|
|
@ -1374,3 +1377,636 @@ class TestSessionPurityInvariant:
|
|||
)
|
||||
assert rejected is SemanticRejectionResult.NOT_DUPLICATE
|
||||
mock_query.assert_not_awaited()
|
||||
|
||||
|
||||
class TestCreateDocumentsConcurrency:
|
||||
"""Concurrent same-collection reinforcements lock rows in id order."""
|
||||
|
||||
N_DOCS: int = 20
|
||||
N_ROUNDS: int = 5
|
||||
|
||||
async def _setup(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
test_workspace: models.Workspace,
|
||||
test_peer: models.Peer,
|
||||
) -> tuple[models.Peer, models.Session]:
|
||||
"""Create an observed peer, session, and collection, committed so
|
||||
they are visible to independent concurrent sessions."""
|
||||
test_peer2 = models.Peer(
|
||||
name=str(generate_nanoid()), workspace_name=test_workspace.name
|
||||
)
|
||||
test_session = models.Session(
|
||||
name=str(generate_nanoid()), workspace_name=test_workspace.name
|
||||
)
|
||||
db_session.add_all([test_peer2, test_session])
|
||||
await db_session.flush()
|
||||
collection = models.Collection(
|
||||
workspace_name=test_workspace.name,
|
||||
observer=test_peer.name,
|
||||
observed=test_peer2.name,
|
||||
)
|
||||
db_session.add(collection)
|
||||
await db_session.commit()
|
||||
return test_peer2, test_session
|
||||
|
||||
def _batch(self, session_name: str) -> list[schemas.DocumentCreate]:
|
||||
return [
|
||||
schemas.DocumentCreate(
|
||||
content=f"user fact number {i}",
|
||||
embedding=[0.1] * 1536,
|
||||
session_name=session_name,
|
||||
metadata=schemas.DocumentMetadata(
|
||||
message_ids=[i],
|
||||
message_created_at="2026-01-01T00:00:00Z",
|
||||
),
|
||||
)
|
||||
for i in range(self.N_DOCS)
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def _chain(exc: BaseException) -> str:
|
||||
parts: list[str] = []
|
||||
seen: set[int] = set()
|
||||
e: BaseException | None = exc
|
||||
while e is not None and id(e) not in seen:
|
||||
seen.add(id(e))
|
||||
parts.append(f"{type(e).__name__}: {e}")
|
||||
e = e.__cause__ or e.__context__
|
||||
return " <- ".join(parts)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_reinforcement_does_not_deadlock(
|
||||
self,
|
||||
db_engine: "AsyncEngine",
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
"""Opposing-order batches on one collection must not deadlock."""
|
||||
test_workspace, test_peer = sample_data
|
||||
test_peer2, test_session = await self._setup(
|
||||
db_session, test_workspace, test_peer
|
||||
)
|
||||
|
||||
# Seed the rows both writers will reinforce.
|
||||
await crud.create_documents(
|
||||
db_session,
|
||||
self._batch(test_session.name),
|
||||
workspace_name=test_workspace.name,
|
||||
observer=test_peer.name,
|
||||
observed=test_peer2.name,
|
||||
)
|
||||
|
||||
session_factory = async_sessionmaker(bind=db_engine, expire_on_commit=False)
|
||||
|
||||
for round_num in range(self.N_ROUNDS):
|
||||
forward = self._batch(test_session.name)
|
||||
backward = list(reversed(self._batch(test_session.name)))
|
||||
|
||||
async def _run(batch: list[schemas.DocumentCreate]) -> None:
|
||||
async with session_factory() as db:
|
||||
await crud.create_documents(
|
||||
db,
|
||||
batch,
|
||||
workspace_name=test_workspace.name,
|
||||
observer=test_peer.name,
|
||||
observed=test_peer2.name,
|
||||
)
|
||||
|
||||
results = await asyncio.gather(
|
||||
_run(forward), _run(backward), return_exceptions=True
|
||||
)
|
||||
errors = [r for r in results if isinstance(r, BaseException)]
|
||||
assert not errors, (
|
||||
f"round {round_num}: concurrent create_documents failed: "
|
||||
+ "; ".join(self._chain(e) for e in errors)
|
||||
)
|
||||
|
||||
# Every round reinforced the same rows: 1 seed + 2 per round.
|
||||
docs = (
|
||||
(
|
||||
await db_session.execute(
|
||||
select(models.Document).where(
|
||||
models.Document.workspace_name == test_workspace.name,
|
||||
models.Document.observer == test_peer.name,
|
||||
models.Document.observed == test_peer2.name,
|
||||
models.Document.deleted_at.is_(None),
|
||||
)
|
||||
)
|
||||
)
|
||||
.scalars()
|
||||
.all()
|
||||
)
|
||||
assert len(docs) == self.N_DOCS
|
||||
assert all(d.times_derived == 1 + 2 * self.N_ROUNDS for d in docs)
|
||||
|
||||
|
||||
class TestCreateDocumentsErrorHandling:
|
||||
"""A dead transaction aborts the batch; per-document failures skip one document."""
|
||||
|
||||
async def _setup(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
test_workspace: models.Workspace,
|
||||
test_peer: models.Peer,
|
||||
) -> tuple[models.Peer, models.Session]:
|
||||
test_peer2 = models.Peer(
|
||||
name=str(generate_nanoid()), workspace_name=test_workspace.name
|
||||
)
|
||||
test_session = models.Session(
|
||||
name=str(generate_nanoid()), workspace_name=test_workspace.name
|
||||
)
|
||||
db_session.add_all([test_peer2, test_session])
|
||||
await db_session.flush()
|
||||
collection = models.Collection(
|
||||
workspace_name=test_workspace.name,
|
||||
observer=test_peer.name,
|
||||
observed=test_peer2.name,
|
||||
)
|
||||
db_session.add(collection)
|
||||
await db_session.commit()
|
||||
return test_peer2, test_session
|
||||
|
||||
def _doc(self, content: str, session_name: str) -> schemas.DocumentCreate:
|
||||
return schemas.DocumentCreate(
|
||||
content=content,
|
||||
embedding=[0.1] * 1536,
|
||||
session_name=session_name,
|
||||
metadata=schemas.DocumentMetadata(
|
||||
message_ids=[1],
|
||||
message_created_at="2026-01-01T00:00:00Z",
|
||||
),
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_db_error_on_row_update_flush_aborts_batch(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
"""A DB error while applying row updates raises and commits nothing."""
|
||||
test_workspace, test_peer = sample_data
|
||||
test_peer2, test_session = await self._setup(
|
||||
db_session, test_workspace, test_peer
|
||||
)
|
||||
# Plain strings: the rollback below expires ORM objects in the session.
|
||||
workspace_name = test_workspace.name
|
||||
observer = test_peer.name
|
||||
observed = test_peer2.name
|
||||
session_name = test_session.name
|
||||
|
||||
await crud.create_documents(
|
||||
db_session,
|
||||
[self._doc("existing fact", session_name)],
|
||||
workspace_name=workspace_name,
|
||||
observer=observer,
|
||||
observed=observed,
|
||||
)
|
||||
|
||||
class FakePGError(Exception):
|
||||
sqlstate: str = "40P01"
|
||||
|
||||
deadlock = OperationalError("UPDATE documents", {}, FakePGError())
|
||||
with (
|
||||
patch.object(db_session, "flush", AsyncMock(side_effect=deadlock)),
|
||||
pytest.raises(OperationalError),
|
||||
):
|
||||
await crud.create_documents(
|
||||
db_session,
|
||||
[
|
||||
self._doc("existing fact", session_name),
|
||||
self._doc("a brand new fact", session_name),
|
||||
],
|
||||
workspace_name=workspace_name,
|
||||
observer=observer,
|
||||
observed=observed,
|
||||
)
|
||||
|
||||
docs = (
|
||||
(
|
||||
await db_session.execute(
|
||||
select(models.Document).where(
|
||||
models.Document.workspace_name == workspace_name,
|
||||
models.Document.observer == observer,
|
||||
models.Document.observed == observed,
|
||||
)
|
||||
)
|
||||
)
|
||||
.scalars()
|
||||
.all()
|
||||
)
|
||||
assert [d.content for d in docs] == ["existing fact"]
|
||||
assert docs[0].times_derived == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_db_error_in_loop_aborts_batch(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
"""A DB error during per-document classification raises and commits nothing."""
|
||||
test_workspace, test_peer = sample_data
|
||||
test_peer2, test_session = await self._setup(
|
||||
db_session, test_workspace, test_peer
|
||||
)
|
||||
workspace_name = test_workspace.name
|
||||
observer = test_peer.name
|
||||
observed = test_peer2.name
|
||||
session_name = test_session.name
|
||||
|
||||
class FakePGError(Exception):
|
||||
sqlstate: str = "40P01"
|
||||
|
||||
deadlock = OperationalError("SELECT documents", {}, FakePGError())
|
||||
with (
|
||||
patch(
|
||||
"src.crud.document._semantic_dup_decision",
|
||||
AsyncMock(side_effect=deadlock),
|
||||
),
|
||||
pytest.raises(OperationalError),
|
||||
):
|
||||
await crud.create_documents(
|
||||
db_session,
|
||||
[
|
||||
self._doc("a brand new fact", session_name),
|
||||
self._doc("another new fact", session_name),
|
||||
],
|
||||
workspace_name=workspace_name,
|
||||
observer=observer,
|
||||
observed=observed,
|
||||
deduplicate=True,
|
||||
)
|
||||
|
||||
docs = (
|
||||
(
|
||||
await db_session.execute(
|
||||
select(models.Document).where(
|
||||
models.Document.workspace_name == workspace_name,
|
||||
models.Document.observer == observer,
|
||||
models.Document.observed == observed,
|
||||
)
|
||||
)
|
||||
)
|
||||
.scalars()
|
||||
.all()
|
||||
)
|
||||
assert docs == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_per_document_error_still_skips_only_that_document(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
"""Non-DB per-document failures keep their skip semantics."""
|
||||
test_workspace, test_peer = sample_data
|
||||
test_peer2, test_session = await self._setup(
|
||||
db_session, test_workspace, test_peer
|
||||
)
|
||||
|
||||
from src.crud import document as document_module
|
||||
|
||||
real_dedup_key = document_module._dedup_key # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
def flaky_dedup_key(
|
||||
content: str, level: str, session_name: str | None
|
||||
) -> tuple[str, str, str | None]:
|
||||
if content == "poison":
|
||||
raise ValueError("bad content")
|
||||
return real_dedup_key(content, level, session_name)
|
||||
|
||||
with patch.object(document_module, "_dedup_key", flaky_dedup_key):
|
||||
result = await crud.create_documents(
|
||||
db_session,
|
||||
[
|
||||
self._doc("good fact one", test_session.name),
|
||||
self._doc("poison", test_session.name),
|
||||
self._doc("good fact two", test_session.name),
|
||||
],
|
||||
workspace_name=test_workspace.name,
|
||||
observer=test_peer.name,
|
||||
observed=test_peer2.name,
|
||||
)
|
||||
|
||||
assert sorted(d.content for d in result.created_documents) == [
|
||||
"good fact one",
|
||||
"good fact two",
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_embedding_skips_semantic_without_embed(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""Empty embeddings must not trigger embed() under an open session."""
|
||||
from src.config import settings
|
||||
|
||||
test_workspace, test_peer = sample_data
|
||||
test_peer2, test_session = await self._setup(
|
||||
db_session, test_workspace, test_peer
|
||||
)
|
||||
monkeypatch.setattr(settings.VECTOR_STORE, "TYPE", "pgvector")
|
||||
monkeypatch.setattr(settings.VECTOR_STORE, "MIGRATED", True)
|
||||
|
||||
empty = self._doc("fact without vector", test_session.name)
|
||||
empty.embedding = []
|
||||
|
||||
with patch(
|
||||
"src.crud.document.embedding_client.embed",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_embed:
|
||||
result = await crud.create_documents(
|
||||
db_session,
|
||||
[empty],
|
||||
workspace_name=test_workspace.name,
|
||||
observer=test_peer.name,
|
||||
observed=test_peer2.name,
|
||||
deduplicate=True,
|
||||
)
|
||||
|
||||
assert len(result.created_documents) == 1
|
||||
mock_embed.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stale_reinforce_target_falls_back_to_insert(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
"""If a reinforce target vanishes under lock, insert the incoming doc."""
|
||||
from src.crud import document as document_module
|
||||
|
||||
test_workspace, test_peer = sample_data
|
||||
test_peer2, test_session = await self._setup(
|
||||
db_session, test_workspace, test_peer
|
||||
)
|
||||
workspace_name = test_workspace.name
|
||||
observer = test_peer.name
|
||||
observed = test_peer2.name
|
||||
session_name = test_session.name
|
||||
|
||||
seeded = await crud.create_documents(
|
||||
db_session,
|
||||
[self._doc("shared fact", session_name)],
|
||||
workspace_name=workspace_name,
|
||||
observer=observer,
|
||||
observed=observed,
|
||||
)
|
||||
assert len(seeded.created_documents) == 1
|
||||
|
||||
existing = (
|
||||
await db_session.execute(
|
||||
select(models.Document).where(
|
||||
models.Document.workspace_name == workspace_name,
|
||||
models.Document.observer == observer,
|
||||
models.Document.observed == observed,
|
||||
models.Document.deleted_at.is_(None),
|
||||
)
|
||||
)
|
||||
).scalar_one()
|
||||
|
||||
real_apply = document_module._apply_document_row_updates # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
async def delete_then_apply(*args: Any, **kwargs: Any) -> Any:
|
||||
existing.deleted_at = datetime.datetime.now(datetime.UTC)
|
||||
await db_session.flush()
|
||||
return await real_apply(*args, **kwargs)
|
||||
|
||||
with patch.object(
|
||||
document_module,
|
||||
"_apply_document_row_updates",
|
||||
side_effect=delete_then_apply,
|
||||
):
|
||||
result = await crud.create_documents(
|
||||
db_session,
|
||||
[self._doc("shared fact", session_name)],
|
||||
workspace_name=workspace_name,
|
||||
observer=observer,
|
||||
observed=observed,
|
||||
)
|
||||
|
||||
assert len(result.created_documents) == 1
|
||||
live = (
|
||||
(
|
||||
await db_session.execute(
|
||||
select(models.Document).where(
|
||||
models.Document.workspace_name == workspace_name,
|
||||
models.Document.observer == observer,
|
||||
models.Document.observed == observed,
|
||||
models.Document.deleted_at.is_(None),
|
||||
)
|
||||
)
|
||||
)
|
||||
.scalars()
|
||||
.all()
|
||||
)
|
||||
assert len(live) == 1
|
||||
assert live[0].id != existing.id
|
||||
assert live[0].content == "shared fact"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_same_batch_replace_then_reinforce_does_not_resurrect(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
"""A reinforce after a same-batch replace must not insert the inferior copy."""
|
||||
from src.crud import document as document_module
|
||||
|
||||
test_workspace, test_peer = sample_data
|
||||
test_peer2, test_session = await self._setup(
|
||||
db_session, test_workspace, test_peer
|
||||
)
|
||||
workspace_name = test_workspace.name
|
||||
observer = test_peer.name
|
||||
observed = test_peer2.name
|
||||
session_name = test_session.name
|
||||
|
||||
await crud.create_documents(
|
||||
db_session,
|
||||
[self._doc("shared fact", session_name)],
|
||||
workspace_name=workspace_name,
|
||||
observer=observer,
|
||||
observed=observed,
|
||||
)
|
||||
existing = (
|
||||
await db_session.execute(
|
||||
select(models.Document).where(
|
||||
models.Document.workspace_name == workspace_name,
|
||||
models.Document.observer == observer,
|
||||
models.Document.observed == observed,
|
||||
models.Document.deleted_at.is_(None),
|
||||
)
|
||||
)
|
||||
).scalar_one()
|
||||
|
||||
fallback = self._doc("shared fact", session_name)
|
||||
ops = [
|
||||
document_module._DocumentRowOp("replace", existing.id), # pyright: ignore[reportPrivateUsage]
|
||||
document_module._DocumentRowOp( # pyright: ignore[reportPrivateUsage]
|
||||
"reinforce",
|
||||
existing.id,
|
||||
fallback_document=fallback,
|
||||
),
|
||||
]
|
||||
fallbacks = await document_module._apply_document_row_updates( # pyright: ignore[reportPrivateUsage]
|
||||
db_session,
|
||||
ops,
|
||||
workspace_name=workspace_name,
|
||||
observer=observer,
|
||||
observed=observed,
|
||||
)
|
||||
assert fallbacks == []
|
||||
await db_session.commit()
|
||||
live = (
|
||||
(
|
||||
await db_session.execute(
|
||||
select(models.Document).where(
|
||||
models.Document.workspace_name == workspace_name,
|
||||
models.Document.observer == observer,
|
||||
models.Document.observed == observed,
|
||||
models.Document.deleted_at.is_(None),
|
||||
)
|
||||
)
|
||||
)
|
||||
.scalars()
|
||||
.all()
|
||||
)
|
||||
assert live == []
|
||||
|
||||
|
||||
class TestExternalCandidateHoist:
|
||||
"""External-store dup candidates resolve before the first DB statement."""
|
||||
|
||||
async def _setup(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
test_workspace: models.Workspace,
|
||||
test_peer: models.Peer,
|
||||
) -> tuple[models.Peer, models.Session]:
|
||||
observed_peer = models.Peer(
|
||||
name=str(generate_nanoid()), workspace_name=test_workspace.name
|
||||
)
|
||||
test_session = models.Session(
|
||||
name=str(generate_nanoid()), workspace_name=test_workspace.name
|
||||
)
|
||||
db_session.add_all([observed_peer, test_session])
|
||||
await db_session.flush()
|
||||
db_session.add(
|
||||
models.Collection(
|
||||
workspace_name=test_workspace.name,
|
||||
observer=test_peer.name,
|
||||
observed=observed_peer.name,
|
||||
)
|
||||
)
|
||||
await db_session.commit()
|
||||
return observed_peer, test_session
|
||||
|
||||
def _doc(self, content: str, session_name: str) -> schemas.DocumentCreate:
|
||||
return schemas.DocumentCreate(
|
||||
content=content,
|
||||
embedding=[0.1] * 1536,
|
||||
session_name=session_name,
|
||||
metadata=schemas.DocumentMetadata(
|
||||
message_ids=[1],
|
||||
message_created_at="2026-01-01T00:00:00Z",
|
||||
),
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_external_candidates_resolved_before_db(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
from src.config import settings
|
||||
|
||||
test_workspace, test_peer = sample_data
|
||||
observed_peer, test_session = await self._setup(
|
||||
db_session, test_workspace, test_peer
|
||||
)
|
||||
monkeypatch.setattr(settings.VECTOR_STORE, "TYPE", "turbopuffer")
|
||||
monkeypatch.setattr(settings.VECTOR_STORE, "MIGRATED", True)
|
||||
|
||||
events: list[str] = []
|
||||
real_execute = db_session.execute
|
||||
|
||||
async def spying_execute(statement: Any, *args: Any, **kwargs: Any) -> Any:
|
||||
events.append("execute")
|
||||
return await real_execute(statement, *args, **kwargs)
|
||||
|
||||
async def fake_resolve(*_args: Any, **_kwargs: Any) -> list[str]:
|
||||
events.append("resolve")
|
||||
return []
|
||||
|
||||
with (
|
||||
patch.object(db_session, "execute", side_effect=spying_execute),
|
||||
patch(
|
||||
"src.crud.document.query_external_vector_document_ids",
|
||||
side_effect=fake_resolve,
|
||||
),
|
||||
patch(
|
||||
"src.crud.document.get_external_vector_store",
|
||||
return_value=None,
|
||||
),
|
||||
):
|
||||
result = await crud.create_documents(
|
||||
db_session,
|
||||
[
|
||||
self._doc("fact one", test_session.name),
|
||||
self._doc("fact two", test_session.name),
|
||||
],
|
||||
workspace_name=test_workspace.name,
|
||||
observer=test_peer.name,
|
||||
observed=observed_peer.name,
|
||||
deduplicate=True,
|
||||
)
|
||||
|
||||
assert len(result.created_documents) == 2
|
||||
assert events[:2] == ["resolve", "resolve"]
|
||||
assert "execute" in events
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_failure_skips_semantic_without_query_documents(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
from src.config import settings
|
||||
|
||||
test_workspace, test_peer = sample_data
|
||||
observed_peer, test_session = await self._setup(
|
||||
db_session, test_workspace, test_peer
|
||||
)
|
||||
monkeypatch.setattr(settings.VECTOR_STORE, "TYPE", "turbopuffer")
|
||||
monkeypatch.setattr(settings.VECTOR_STORE, "MIGRATED", True)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"src.crud.document.query_external_vector_document_ids",
|
||||
side_effect=RuntimeError("store down"),
|
||||
),
|
||||
patch(
|
||||
"src.crud.document.get_external_vector_store",
|
||||
return_value=None,
|
||||
),
|
||||
patch(
|
||||
"src.crud.document.query_documents",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_query,
|
||||
):
|
||||
result = await crud.create_documents(
|
||||
db_session,
|
||||
[self._doc("fact one", test_session.name)],
|
||||
workspace_name=test_workspace.name,
|
||||
observer=test_peer.name,
|
||||
observed=observed_peer.name,
|
||||
deduplicate=True,
|
||||
)
|
||||
|
||||
assert len(result.created_documents) == 1
|
||||
mock_query.assert_not_awaited()
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
from contextlib import asynccontextmanager
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -196,7 +196,7 @@ class TestRepresentationManagerSoftDelete:
|
|||
db_session, test_workspace, test_peer
|
||||
)
|
||||
|
||||
base = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||||
base = datetime(2026, 1, 1, tzinfo=UTC)
|
||||
# Three conclusions, all reinforced once, inserted oldest-first.
|
||||
for i in range(3):
|
||||
db_session.add(
|
||||
|
|
@ -484,13 +484,13 @@ class TestRepresentationManagerSave:
|
|||
explicit=[
|
||||
ExplicitObservation(
|
||||
content=" ",
|
||||
created_at=datetime.now(timezone.utc),
|
||||
created_at=datetime.now(UTC),
|
||||
message_ids=[1],
|
||||
session_name="session",
|
||||
),
|
||||
ExplicitObservation(
|
||||
content=" useful observation ",
|
||||
created_at=datetime.now(timezone.utc),
|
||||
created_at=datetime.now(UTC),
|
||||
message_ids=[1],
|
||||
session_name="session",
|
||||
),
|
||||
|
|
@ -515,7 +515,7 @@ class TestRepresentationManagerSave:
|
|||
representation,
|
||||
message_ids=[1],
|
||||
session_name="session",
|
||||
message_created_at=datetime.now(timezone.utc),
|
||||
message_created_at=datetime.now(UTC),
|
||||
message_level_configuration=_resolved_config(),
|
||||
)
|
||||
|
||||
|
|
@ -540,7 +540,7 @@ class TestRepresentationManagerSave:
|
|||
conclusion=" ",
|
||||
premises=["premise a"],
|
||||
source_ids=["doc-a"],
|
||||
created_at=datetime.now(timezone.utc),
|
||||
created_at=datetime.now(UTC),
|
||||
message_ids=[1],
|
||||
session_name="session",
|
||||
),
|
||||
|
|
@ -548,7 +548,7 @@ class TestRepresentationManagerSave:
|
|||
conclusion=" inferred conclusion ",
|
||||
premises=["premise b"],
|
||||
source_ids=["doc-b"],
|
||||
created_at=datetime.now(timezone.utc),
|
||||
created_at=datetime.now(UTC),
|
||||
message_ids=[1],
|
||||
session_name="session",
|
||||
),
|
||||
|
|
@ -573,7 +573,7 @@ class TestRepresentationManagerSave:
|
|||
representation,
|
||||
message_ids=[1],
|
||||
session_name="session",
|
||||
message_created_at=datetime.now(timezone.utc),
|
||||
message_created_at=datetime.now(UTC),
|
||||
message_level_configuration=_resolved_config(),
|
||||
)
|
||||
|
||||
|
|
@ -597,13 +597,13 @@ class TestRepresentationManagerSave:
|
|||
explicit=[
|
||||
ExplicitObservation(
|
||||
content="",
|
||||
created_at=datetime.now(timezone.utc),
|
||||
created_at=datetime.now(UTC),
|
||||
message_ids=[1],
|
||||
session_name="session",
|
||||
),
|
||||
ExplicitObservation(
|
||||
content="\n\t ",
|
||||
created_at=datetime.now(timezone.utc),
|
||||
created_at=datetime.now(UTC),
|
||||
message_ids=[1],
|
||||
session_name="session",
|
||||
),
|
||||
|
|
@ -626,7 +626,124 @@ class TestRepresentationManagerSave:
|
|||
representation,
|
||||
message_ids=[1],
|
||||
session_name="session",
|
||||
message_created_at=datetime.now(timezone.utc),
|
||||
message_created_at=datetime.now(UTC),
|
||||
message_level_configuration=_resolved_config(),
|
||||
)
|
||||
|
||||
assert len(saved.created_documents) == 0
|
||||
mock_embed.assert_not_awaited()
|
||||
mock_save.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_representation_strips_nul_bytes(self):
|
||||
"""Models emit \\u0000 escapes when transcribing shell output or Windows
|
||||
paths, and Postgres rejects NUL in text columns. The stripped text must
|
||||
be what gets embedded as well as what gets stored."""
|
||||
manager = RepresentationManager(
|
||||
"workspace",
|
||||
observer="observer",
|
||||
observed="observed",
|
||||
)
|
||||
representation = Representation(
|
||||
explicit=[
|
||||
ExplicitObservation(
|
||||
content="ran 'cat /proc/1/environ | tr '\x00' '\\n''",
|
||||
created_at=datetime.now(UTC),
|
||||
message_ids=[1],
|
||||
session_name="session",
|
||||
),
|
||||
],
|
||||
deductive=[
|
||||
DeductiveObservation(
|
||||
conclusion="the key is at c:\\\x00users\\amal",
|
||||
premises=["saw c:\\\x00users in the prompt"],
|
||||
created_at=datetime.now(UTC),
|
||||
message_ids=[1],
|
||||
session_name="session",
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
with (
|
||||
patch("src.crud.representation.tracked_db", _fake_tracked_db),
|
||||
patch(
|
||||
"src.crud.representation.embedding_client.simple_batch_embed",
|
||||
new=AsyncMock(return_value=[[0.1], [0.2]]),
|
||||
) as mock_embed,
|
||||
patch.object(
|
||||
manager,
|
||||
"_save_representation_internal",
|
||||
new=AsyncMock(
|
||||
return_value=CreateDocumentsResult(created_documents=[MagicMock()])
|
||||
),
|
||||
) as mock_save,
|
||||
):
|
||||
await manager.save_representation(
|
||||
representation,
|
||||
message_ids=[1],
|
||||
session_name="session",
|
||||
message_created_at=datetime.now(UTC),
|
||||
message_level_configuration=_resolved_config(),
|
||||
)
|
||||
|
||||
# Deductive observations are embedded ahead of explicit ones.
|
||||
mock_embed.assert_awaited_once_with(
|
||||
[
|
||||
"the key is at c:\\users\\amal",
|
||||
"ran 'cat /proc/1/environ | tr '' '\\n''",
|
||||
],
|
||||
on_oversize="truncate",
|
||||
)
|
||||
|
||||
saved_observations = _saved_observations(mock_save)
|
||||
deductive = next(
|
||||
obs for obs in saved_observations if isinstance(obs, DeductiveObservation)
|
||||
)
|
||||
explicit = next(
|
||||
obs for obs in saved_observations if isinstance(obs, ExplicitObservation)
|
||||
)
|
||||
assert explicit.content == "ran 'cat /proc/1/environ | tr '' '\\n''"
|
||||
assert deductive.conclusion == "the key is at c:\\users\\amal"
|
||||
# premises land in internal_metadata, and jsonb rejects NUL too
|
||||
assert deductive.premises == ["saw c:\\users in the prompt"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_representation_skips_observations_that_are_only_nul(self):
|
||||
"""str.strip() does not remove NUL, so the emptiness check has to run
|
||||
after normalization or an empty document gets written."""
|
||||
manager = RepresentationManager(
|
||||
"workspace",
|
||||
observer="observer",
|
||||
observed="observed",
|
||||
)
|
||||
representation = Representation(
|
||||
explicit=[
|
||||
ExplicitObservation(
|
||||
content="\x00\x00",
|
||||
created_at=datetime.now(UTC),
|
||||
message_ids=[1],
|
||||
session_name="session",
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
with (
|
||||
patch("src.crud.representation.tracked_db", _fake_tracked_db),
|
||||
patch(
|
||||
"src.crud.representation.embedding_client.simple_batch_embed",
|
||||
new=AsyncMock(),
|
||||
) as mock_embed,
|
||||
patch.object(
|
||||
manager,
|
||||
"_save_representation_internal",
|
||||
new=AsyncMock(),
|
||||
) as mock_save,
|
||||
):
|
||||
saved = await manager.save_representation(
|
||||
representation,
|
||||
message_ids=[1],
|
||||
session_name="session",
|
||||
message_created_at=datetime.now(UTC),
|
||||
message_level_configuration=_resolved_config(),
|
||||
)
|
||||
|
||||
|
|
@ -646,7 +763,7 @@ class TestRepresentationManagerSave:
|
|||
explicit=[
|
||||
ExplicitObservation(
|
||||
content="short fact",
|
||||
created_at=datetime.now(timezone.utc),
|
||||
created_at=datetime.now(UTC),
|
||||
message_ids=[1],
|
||||
session_name="session",
|
||||
)
|
||||
|
|
@ -656,7 +773,7 @@ class TestRepresentationManagerSave:
|
|||
conclusion="inferred fact",
|
||||
premises=["premise"],
|
||||
source_ids=["doc-a"],
|
||||
created_at=datetime.now(timezone.utc),
|
||||
created_at=datetime.now(UTC),
|
||||
message_ids=[1],
|
||||
session_name="session",
|
||||
)
|
||||
|
|
@ -681,7 +798,7 @@ class TestRepresentationManagerSave:
|
|||
representation,
|
||||
message_ids=[1],
|
||||
session_name="session",
|
||||
message_created_at=datetime.now(timezone.utc),
|
||||
message_created_at=datetime.now(UTC),
|
||||
message_level_configuration=_resolved_config(),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
import signal
|
||||
from datetime import datetime, timezone
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
|
|
@ -32,7 +32,7 @@ class TestDeriverProcessing:
|
|||
peer_name="alice",
|
||||
content="hello",
|
||||
token_count=5,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
created_at=datetime.now(UTC),
|
||||
)
|
||||
configuration = Mock()
|
||||
configuration.reasoning.enabled = True
|
||||
|
|
@ -82,7 +82,7 @@ class TestDeriverProcessing:
|
|||
peer_name="alice",
|
||||
content="hello",
|
||||
token_count=5,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
created_at=datetime.now(UTC),
|
||||
)
|
||||
configuration = Mock()
|
||||
configuration.reasoning.enabled = True
|
||||
|
|
@ -136,7 +136,7 @@ class TestDeriverProcessing:
|
|||
peer_name="alice",
|
||||
content="hello",
|
||||
token_count=5,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
created_at=datetime.now(UTC),
|
||||
)
|
||||
configuration = Mock()
|
||||
configuration.reasoning.enabled = True
|
||||
|
|
@ -182,6 +182,61 @@ class TestDeriverProcessing:
|
|||
assert event.observer_count == 1
|
||||
assert event.failed_observer_count == 1
|
||||
|
||||
async def test_retryable_observer_save_reraises_after_telemetry(self):
|
||||
"""A deadlock on one observer must propagate so the queue can retry."""
|
||||
from sqlalchemy.exc import OperationalError
|
||||
|
||||
class FakePGError(Exception):
|
||||
sqlstate: str = "40P01"
|
||||
|
||||
deadlock = OperationalError("UPDATE documents", {}, FakePGError())
|
||||
message = Mock(
|
||||
id=1,
|
||||
public_id="msg_1",
|
||||
session_name="session-1",
|
||||
workspace_name="workspace-1",
|
||||
peer_name="alice",
|
||||
content="hello",
|
||||
token_count=5,
|
||||
created_at=datetime.now(UTC),
|
||||
)
|
||||
configuration = Mock()
|
||||
configuration.reasoning.enabled = True
|
||||
|
||||
mock_response = HonchoLLMCallResponse(
|
||||
content=PromptRepresentation(
|
||||
explicit=[
|
||||
ExplicitObservationBase(content="The user has a dog named Rover")
|
||||
]
|
||||
),
|
||||
input_tokens=10,
|
||||
output_tokens=5,
|
||||
finish_reasons=["STOP"],
|
||||
)
|
||||
partial_save = AsyncMock(side_effect=[crud.CreateDocumentsResult(), deadlock])
|
||||
emitted: list[Any] = []
|
||||
with (
|
||||
patch(
|
||||
"src.deriver.deriver.honcho_llm_call",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_response,
|
||||
),
|
||||
patch.object(RepresentationManager, "save_representation", partial_save),
|
||||
patch("src.deriver.deriver.emit", side_effect=emitted.append),
|
||||
pytest.raises(OperationalError),
|
||||
):
|
||||
await process_representation_tasks_batch(
|
||||
messages=[message],
|
||||
message_level_configuration=configuration,
|
||||
observers=["bob", "carol"],
|
||||
observed="alice",
|
||||
queue_item_message_ids=[1],
|
||||
)
|
||||
|
||||
assert emitted, "expected telemetry to be emitted before the raised failure"
|
||||
assert emitted[-1].observer_count == 1
|
||||
assert emitted[-1].failed_observer_count == 1
|
||||
|
||||
async def test_process_representation_tasks_batch_passes_custom_instructions_into_prompt(
|
||||
self,
|
||||
) -> None:
|
||||
|
|
@ -193,7 +248,7 @@ class TestDeriverProcessing:
|
|||
peer_name="alice",
|
||||
content="hello",
|
||||
token_count=5,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
created_at=datetime.now(UTC),
|
||||
)
|
||||
configuration = Mock()
|
||||
configuration.reasoning.enabled = True
|
||||
|
|
@ -343,7 +398,7 @@ class TestDeriverProcessing:
|
|||
peer_name="alice",
|
||||
content="hello",
|
||||
token_count=100,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
created_at=datetime.now(UTC),
|
||||
)
|
||||
configuration = Mock()
|
||||
configuration.reasoning.enabled = True
|
||||
|
|
@ -394,7 +449,7 @@ class TestDeriverProcessing:
|
|||
peer_name="alice",
|
||||
content="hello",
|
||||
token_count=5,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
created_at=datetime.now(UTC),
|
||||
)
|
||||
configuration = Mock()
|
||||
configuration.reasoning.enabled = True
|
||||
|
|
@ -443,7 +498,7 @@ class TestDeriverProcessing:
|
|||
peer_name="alice",
|
||||
content="hello",
|
||||
token_count=5,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
created_at=datetime.now(UTC),
|
||||
)
|
||||
configuration = Mock()
|
||||
configuration.reasoning.enabled = True
|
||||
|
|
|
|||
|
|
@ -1,17 +1,21 @@
|
|||
import asyncio
|
||||
from collections.abc import Callable
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from nanoid import generate as generate_nanoid
|
||||
from pydantic import ValidationError
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.exc import OperationalError
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from src import models
|
||||
from src.config import settings
|
||||
from src.deriver.consumer import process_item
|
||||
from src.deriver.queue_manager import QueueManager, WorkerOwnership
|
||||
from src.utils.queue_payload import RETRY_ATTEMPTS_PAYLOAD_KEY, SummaryPayload
|
||||
from src.utils.work_unit import construct_work_unit_key
|
||||
|
||||
|
||||
|
|
@ -1519,7 +1523,7 @@ class TestQueueProcessing:
|
|||
monkeypatch.setattr(
|
||||
settings.DERIVER, "REPRESENTATION_BATCH_MAX_AGE_SECONDS", 1800
|
||||
)
|
||||
old_timestamp = datetime.now(timezone.utc) - timedelta(hours=2)
|
||||
old_timestamp = datetime.now(UTC) - timedelta(hours=2)
|
||||
|
||||
work_unit_key, queue_items = await self._add_representation_work_unit(
|
||||
db_session=db_session,
|
||||
|
|
@ -1552,7 +1556,7 @@ class TestQueueProcessing:
|
|||
) -> None:
|
||||
monkeypatch.setattr(settings.DERIVER, "FLUSH_ENABLED", False)
|
||||
monkeypatch.setattr(settings.DERIVER, "REPRESENTATION_BATCH_MAX_AGE_SECONDS", 0)
|
||||
old_timestamp = datetime.now(timezone.utc) - timedelta(hours=2)
|
||||
old_timestamp = datetime.now(UTC) - timedelta(hours=2)
|
||||
|
||||
work_unit_key, _queue_items = await self._add_representation_work_unit(
|
||||
db_session=db_session,
|
||||
|
|
@ -1602,7 +1606,7 @@ class TestQueueProcessing:
|
|||
monkeypatch.setattr(
|
||||
settings.DERIVER, "REPRESENTATION_BATCH_MAX_AGE_SECONDS", 1800
|
||||
)
|
||||
now = datetime.now(timezone.utc)
|
||||
now = datetime.now(UTC)
|
||||
|
||||
work_unit_key, _queue_items = await self._add_representation_work_unit(
|
||||
db_session=db_session,
|
||||
|
|
@ -1628,7 +1632,7 @@ class TestQueueProcessing:
|
|||
monkeypatch.setattr(
|
||||
settings.DERIVER, "REPRESENTATION_BATCH_MAX_AGE_SECONDS", 1800
|
||||
)
|
||||
now = datetime.now(timezone.utc)
|
||||
now = datetime.now(UTC)
|
||||
|
||||
work_unit_key, queue_items = await self._add_representation_work_unit(
|
||||
db_session=db_session,
|
||||
|
|
@ -1874,3 +1878,354 @@ class TestPollingJitter:
|
|||
qm.shutdown_event.set()
|
||||
# A shutdown already signalled must short-circuit the (long) jitter sleep.
|
||||
await asyncio.wait_for(qm._sleep_startup_jitter(), timeout=1.0) # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestQueueRetry:
|
||||
"""Bounded retry of transient errors in process_work_unit (DEV-1975).
|
||||
|
||||
A transient failure (deadlock, lost connection, provider transport) must
|
||||
leave the batch's queue items unprocessed and release the work unit for
|
||||
re-claim, up to MAX_RETRYABLE_ATTEMPTS per work unit; terminal failures
|
||||
keep today's burn-one-item behavior.
|
||||
"""
|
||||
|
||||
async def _seed_work_unit(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_session_with_peers: tuple[models.Session, list[models.Peer]],
|
||||
create_queue_payload: Callable[..., Any],
|
||||
n_messages: int = 1,
|
||||
) -> tuple[QueueManager, str, str, list[models.QueueItem]]:
|
||||
"""Seed a claimed representation work unit owned by a test worker."""
|
||||
session, peers = sample_session_with_peers
|
||||
peer = peers[0]
|
||||
|
||||
messages: list[models.Message] = []
|
||||
for index in range(n_messages):
|
||||
message = models.Message(
|
||||
session_name=session.name,
|
||||
workspace_name=session.workspace_name,
|
||||
peer_name=peer.name,
|
||||
content=f"Message {index}",
|
||||
token_count=10,
|
||||
seq_in_session=index + 1,
|
||||
)
|
||||
db_session.add(message)
|
||||
messages.append(message)
|
||||
await db_session.commit()
|
||||
for message in messages:
|
||||
await db_session.refresh(message)
|
||||
|
||||
queue_items: list[models.QueueItem] = []
|
||||
work_unit_key = ""
|
||||
for message in messages:
|
||||
payload = create_queue_payload(
|
||||
message=message,
|
||||
task_type="representation",
|
||||
observed=peer.name,
|
||||
observer=peer.name,
|
||||
)
|
||||
work_unit_key = work_unit_key or construct_work_unit_key(
|
||||
session.workspace_name, payload
|
||||
)
|
||||
queue_item = models.QueueItem(
|
||||
session_id=session.id,
|
||||
task_type="representation",
|
||||
work_unit_key=work_unit_key,
|
||||
payload=payload,
|
||||
processed=False,
|
||||
workspace_name=session.workspace_name,
|
||||
message_id=message.id,
|
||||
)
|
||||
db_session.add(queue_item)
|
||||
queue_items.append(queue_item)
|
||||
await db_session.commit()
|
||||
for queue_item in queue_items:
|
||||
await db_session.refresh(queue_item)
|
||||
|
||||
qm = QueueManager()
|
||||
worker_id = "test_worker"
|
||||
claimed_units = await qm.claim_work_units(db_session, [work_unit_key])
|
||||
qm.worker_ownership[worker_id] = WorkerOwnership(
|
||||
work_unit_key=work_unit_key, aqs_id=claimed_units[work_unit_key]
|
||||
)
|
||||
await db_session.commit()
|
||||
return qm, work_unit_key, worker_id, queue_items
|
||||
|
||||
@staticmethod
|
||||
def _retryable_error() -> OperationalError:
|
||||
class FakePGError(Exception):
|
||||
sqlstate: str = "40P01"
|
||||
|
||||
return OperationalError("UPDATE documents", {}, FakePGError())
|
||||
|
||||
async def _fetch_items(
|
||||
self, db_session: AsyncSession, work_unit_key: str
|
||||
) -> list[models.QueueItem]:
|
||||
db_session.expire_all()
|
||||
return list(
|
||||
(
|
||||
await db_session.execute(
|
||||
select(models.QueueItem)
|
||||
.where(models.QueueItem.work_unit_key == work_unit_key)
|
||||
.order_by(models.QueueItem.id)
|
||||
)
|
||||
)
|
||||
.scalars()
|
||||
.all()
|
||||
)
|
||||
|
||||
async def _aqs_rows(self, db_session: AsyncSession, work_unit_key: str) -> int:
|
||||
return len(
|
||||
(
|
||||
await db_session.execute(
|
||||
select(models.ActiveQueueSession).where(
|
||||
models.ActiveQueueSession.work_unit_key == work_unit_key
|
||||
)
|
||||
)
|
||||
)
|
||||
.scalars()
|
||||
.all()
|
||||
)
|
||||
|
||||
async def _retry_attempts_on_items(
|
||||
self, db_session: AsyncSession, work_unit_key: str
|
||||
) -> int | None:
|
||||
items = await self._fetch_items(db_session, work_unit_key)
|
||||
unprocessed = [item for item in items if not item.processed]
|
||||
if not unprocessed:
|
||||
return None
|
||||
raw = (unprocessed[0].payload or {}).get("_retry_attempts")
|
||||
return None if raw is None else int(raw)
|
||||
|
||||
async def test_retryable_error_leaves_items_unprocessed(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_session_with_peers: tuple[models.Session, list[models.Peer]],
|
||||
create_queue_payload: Callable[..., Any],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""A transient error stops the work unit after ONE batch fetch (no
|
||||
tight loop), leaves items unprocessed with no error, and releases
|
||||
the ActiveQueueSession row."""
|
||||
monkeypatch.setattr("src.deriver.queue_manager.RETRY_BACKOFF_SECONDS", 0.0)
|
||||
qm, work_unit_key, worker_id, _ = await self._seed_work_unit(
|
||||
db_session, sample_session_with_peers, create_queue_payload, n_messages=2
|
||||
)
|
||||
initial_semaphore_value = qm.semaphore._value
|
||||
|
||||
batch_fetches = 0
|
||||
original_get_batch = qm.get_queue_item_batch
|
||||
|
||||
async def counting_get_batch(*args: Any, **kwargs: Any) -> Any:
|
||||
nonlocal batch_fetches
|
||||
batch_fetches += 1
|
||||
return await original_get_batch(*args, **kwargs)
|
||||
|
||||
with (
|
||||
patch.object(qm, "get_queue_item_batch", side_effect=counting_get_batch),
|
||||
patch(
|
||||
"src.deriver.queue_manager.process_representation_batch",
|
||||
side_effect=self._retryable_error(),
|
||||
),
|
||||
):
|
||||
await qm.process_work_unit(work_unit_key, worker_id)
|
||||
|
||||
assert batch_fetches == 1
|
||||
items = await self._fetch_items(db_session, work_unit_key)
|
||||
assert all(not item.processed for item in items)
|
||||
assert all(item.error is None for item in items)
|
||||
assert await self._aqs_rows(db_session, work_unit_key) == 0
|
||||
assert await self._retry_attempts_on_items(db_session, work_unit_key) == 1
|
||||
assert qm.semaphore._value == initial_semaphore_value
|
||||
|
||||
async def test_retry_exhaustion_is_terminal(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_session_with_peers: tuple[models.Session, list[models.Peer]],
|
||||
create_queue_payload: Callable[..., Any],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""At the attempt cap a transient error burns the first item exactly
|
||||
like today's terminal path and clears the counter."""
|
||||
from src.deriver.queue_manager import MAX_RETRYABLE_ATTEMPTS
|
||||
|
||||
monkeypatch.setattr("src.deriver.queue_manager.RETRY_BACKOFF_SECONDS", 0.0)
|
||||
qm, work_unit_key, worker_id, _ = await self._seed_work_unit(
|
||||
db_session, sample_session_with_peers, create_queue_payload
|
||||
)
|
||||
await qm._set_work_unit_retry_attempts( # pyright: ignore[reportPrivateUsage]
|
||||
work_unit_key, MAX_RETRYABLE_ATTEMPTS - 1
|
||||
)
|
||||
|
||||
with patch(
|
||||
"src.deriver.queue_manager.process_representation_batch",
|
||||
side_effect=self._retryable_error(),
|
||||
):
|
||||
await qm.process_work_unit(work_unit_key, worker_id)
|
||||
|
||||
items = await self._fetch_items(db_session, work_unit_key)
|
||||
assert len(items) == 1
|
||||
assert items[0].processed
|
||||
assert items[0].error is not None
|
||||
assert "OperationalError" in items[0].error
|
||||
assert await self._retry_attempts_on_items(db_session, work_unit_key) is None
|
||||
|
||||
async def test_non_retryable_error_burns_immediately(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_session_with_peers: tuple[models.Session, list[models.Peer]],
|
||||
create_queue_payload: Callable[..., Any],
|
||||
) -> None:
|
||||
"""A non-retryable error keeps today's behavior verbatim: the first
|
||||
item is marked errored on the first attempt."""
|
||||
qm, work_unit_key, worker_id, _ = await self._seed_work_unit(
|
||||
db_session, sample_session_with_peers, create_queue_payload
|
||||
)
|
||||
|
||||
with patch(
|
||||
"src.deriver.queue_manager.process_representation_batch",
|
||||
side_effect=ValueError("bad batch"),
|
||||
):
|
||||
await qm.process_work_unit(work_unit_key, worker_id)
|
||||
|
||||
items = await self._fetch_items(db_session, work_unit_key)
|
||||
assert len(items) == 1
|
||||
assert items[0].processed
|
||||
assert items[0].error is not None
|
||||
assert "ValueError" in items[0].error
|
||||
assert await self._retry_attempts_on_items(db_session, work_unit_key) is None
|
||||
|
||||
async def test_counter_cleared_after_success(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_session_with_peers: tuple[models.Session, list[models.Peer]],
|
||||
create_queue_payload: Callable[..., Any],
|
||||
) -> None:
|
||||
"""A success wipes the accumulated attempt count for the work unit."""
|
||||
qm, work_unit_key, worker_id, _ = await self._seed_work_unit(
|
||||
db_session, sample_session_with_peers, create_queue_payload
|
||||
)
|
||||
await qm._set_work_unit_retry_attempts(work_unit_key, 1) # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
async def noop_batch(*_args: Any, **_kwargs: Any) -> None:
|
||||
return None
|
||||
|
||||
with patch(
|
||||
"src.deriver.queue_manager.process_representation_batch",
|
||||
side_effect=noop_batch,
|
||||
):
|
||||
await qm.process_work_unit(work_unit_key, worker_id)
|
||||
|
||||
items = await self._fetch_items(db_session, work_unit_key)
|
||||
assert all(item.processed for item in items)
|
||||
assert all(item.error is None for item in items)
|
||||
# Counter lives on the oldest unprocessed item; once that item is
|
||||
# processed the budget is gone even if the payload key remains.
|
||||
assert await self._retry_attempts_on_items(db_session, work_unit_key) is None
|
||||
|
||||
async def test_retry_budget_survives_reclaim_by_another_manager(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_session_with_peers: tuple[models.Session, list[models.Peer]],
|
||||
create_queue_payload: Callable[..., Any],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""A second QueueManager continues the durable attempt budget."""
|
||||
from src.deriver.queue_manager import MAX_RETRYABLE_ATTEMPTS
|
||||
|
||||
monkeypatch.setattr("src.deriver.queue_manager.RETRY_BACKOFF_SECONDS", 0.0)
|
||||
qm1, work_unit_key, worker_id, _ = await self._seed_work_unit(
|
||||
db_session, sample_session_with_peers, create_queue_payload
|
||||
)
|
||||
|
||||
with patch(
|
||||
"src.deriver.queue_manager.process_representation_batch",
|
||||
side_effect=self._retryable_error(),
|
||||
):
|
||||
await qm1.process_work_unit(work_unit_key, worker_id)
|
||||
|
||||
assert await self._retry_attempts_on_items(db_session, work_unit_key) == 1
|
||||
assert await self._aqs_rows(db_session, work_unit_key) == 0
|
||||
|
||||
# Seed the remaining budget so the next reclaim is the terminal attempt.
|
||||
qm2 = QueueManager()
|
||||
await qm2._set_work_unit_retry_attempts( # pyright: ignore[reportPrivateUsage]
|
||||
work_unit_key, MAX_RETRYABLE_ATTEMPTS - 1
|
||||
)
|
||||
claimed = await qm2.claim_work_units(db_session, [work_unit_key])
|
||||
worker_id_2 = "test_worker_2"
|
||||
qm2.worker_ownership[worker_id_2] = WorkerOwnership(
|
||||
work_unit_key=work_unit_key, aqs_id=claimed[work_unit_key]
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
with patch(
|
||||
"src.deriver.queue_manager.process_representation_batch",
|
||||
side_effect=self._retryable_error(),
|
||||
):
|
||||
await qm2.process_work_unit(work_unit_key, worker_id_2)
|
||||
|
||||
items = await self._fetch_items(db_session, work_unit_key)
|
||||
assert len(items) == 1
|
||||
assert items[0].processed
|
||||
assert items[0].error is not None
|
||||
assert "OperationalError" in items[0].error
|
||||
|
||||
async def test_process_item_strips_retry_counter_before_validation(self) -> None:
|
||||
"""A reclaimed non-representation item must survive its own retry counter.
|
||||
|
||||
The counter is written onto an *unprocessed* item so the budget outlives
|
||||
a work-unit reclaim -- which means the next claim re-reads it. Every
|
||||
payload model sets ``extra="forbid"``, so without the strip in
|
||||
``process_item`` the reclaim raises extra_forbidden -> ValueError ->
|
||||
not retryable -> the item is burned terminally on the very attempt that
|
||||
was supposed to retry it. Representation tasks never hit this: their
|
||||
batch path reads the payload with ``.get()`` instead of validating,
|
||||
which is why the rest of this class cannot catch it.
|
||||
"""
|
||||
raw: dict[str, Any] = {
|
||||
"task_type": "summary",
|
||||
"session_name": "s",
|
||||
"message_seq_in_session": 1,
|
||||
"message_public_id": "msg-public-id",
|
||||
"configuration": {
|
||||
"reasoning": {"enabled": True},
|
||||
"peer_card": {"use": True, "create": True},
|
||||
"summary": {
|
||||
"enabled": True,
|
||||
"messages_per_short_summary": 20,
|
||||
"messages_per_long_summary": 60,
|
||||
},
|
||||
"dream": {"enabled": True},
|
||||
},
|
||||
RETRY_ATTEMPTS_PAYLOAD_KEY: 1,
|
||||
}
|
||||
|
||||
# Pin the premise: the payload model must keep rejecting the key, so
|
||||
# this fails loudly if someone "fixes" the burn with extra="allow"
|
||||
# instead of stripping.
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
SummaryPayload.model_validate(raw)
|
||||
assert any(err["type"] == "extra_forbidden" for err in exc_info.value.errors())
|
||||
|
||||
queue_item = models.QueueItem(
|
||||
task_type="summary",
|
||||
work_unit_key="summary:test-workspace:test-session",
|
||||
payload=raw,
|
||||
processed=False,
|
||||
workspace_name="test-workspace",
|
||||
message_id=1,
|
||||
)
|
||||
|
||||
with patch(
|
||||
"src.deriver.consumer.summarizer.summarize_if_needed",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_summarize:
|
||||
await process_item(queue_item)
|
||||
|
||||
mock_summarize.assert_awaited_once()
|
||||
# The strip must happen on a copy: the counter has to stay on the row so
|
||||
# the budget still advances if this attempt fails again.
|
||||
assert raw[RETRY_ATTEMPTS_PAYLOAD_KEY] == 1
|
||||
|
|
|
|||
|
|
@ -22,15 +22,19 @@ engine (see ``mock_tracked_db_context`` in conftest.py) — a different
|
|||
connection that cannot see another session's uncommitted writes.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import AsyncGenerator
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
from nanoid import generate as generate_nanoid
|
||||
from sqlalchemy import func, select, update
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker
|
||||
|
||||
from src import crud, models
|
||||
from src.deriver import scope_backfill as scope_backfill_mod
|
||||
from src.deriver.scope_backfill import (
|
||||
COPIED_FROM_KEY,
|
||||
process_scope_backfill,
|
||||
|
|
@ -415,6 +419,128 @@ async def test_backfill_skips_a_session_that_left_the_scope(
|
|||
assert session_name not in peer.internal_metadata.get("backfill_status", {})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_copy_chunk_membership_lock_blocks_leave_until_write_commits(
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
db_engine: AsyncEngine,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""A concurrent leave cannot commit between membership check and inserts."""
|
||||
test_workspace, sender = sample_data
|
||||
workspace_name = test_workspace.name
|
||||
scope_name = str(generate_nanoid())
|
||||
scope_peer = await _create_scope_peer(db_session, workspace_name, scope_name)
|
||||
session = await _create_session(db_session, workspace_name)
|
||||
await _join_scope(db_session, workspace_name, session.name, scope_peer.name)
|
||||
await _create_collection(
|
||||
db_session, workspace_name, observer=sender.name, observed=sender.name
|
||||
)
|
||||
await _create_collection(
|
||||
db_session, workspace_name, observer=scope_peer.name, observed=sender.name
|
||||
)
|
||||
source = await _create_document(
|
||||
db_session,
|
||||
workspace_name,
|
||||
observer=sender.name,
|
||||
observed=sender.name,
|
||||
session_name=session.name,
|
||||
content="locked membership fact",
|
||||
)
|
||||
|
||||
factory = async_sessionmaker(bind=db_engine, expire_on_commit=False)
|
||||
leave_finished = asyncio.Event()
|
||||
leave_task_box: dict[str, asyncio.Task[None]] = {}
|
||||
|
||||
async def concurrent_leave() -> None:
|
||||
async with factory() as leave_db:
|
||||
await leave_db.execute(
|
||||
update(models.SessionPeer)
|
||||
.where(
|
||||
models.SessionPeer.workspace_name == workspace_name,
|
||||
models.SessionPeer.session_name == session.name,
|
||||
models.SessionPeer.peer_name == scope_peer.name,
|
||||
models.SessionPeer.left_at.is_(None),
|
||||
)
|
||||
.values(left_at=func.now())
|
||||
)
|
||||
await leave_db.commit()
|
||||
leave_finished.set()
|
||||
|
||||
original_tracked_db = scope_backfill_mod.tracked_db # pyright: ignore[reportPrivateLocalImportUsage]
|
||||
|
||||
@asynccontextmanager
|
||||
async def tracked_db_with_leave_race(
|
||||
operation_name: str | None = None, *, read_only: bool = False
|
||||
) -> AsyncGenerator[AsyncSession]:
|
||||
async with original_tracked_db(operation_name, read_only=read_only) as db:
|
||||
if operation_name == "scope_backfill.write":
|
||||
real_scalar = db.scalar
|
||||
raced = False
|
||||
|
||||
async def scalar_then_race(statement: Any, *args: Any, **kwargs: Any):
|
||||
nonlocal raced
|
||||
result = await real_scalar(statement, *args, **kwargs)
|
||||
if not raced and result is not None:
|
||||
raced = True
|
||||
leave_task_box["task"] = asyncio.create_task(concurrent_leave())
|
||||
# Leave's UPDATE must block on this txn's row lock.
|
||||
for _ in range(50):
|
||||
await asyncio.sleep(0.01)
|
||||
if leave_task_box["task"].done():
|
||||
break
|
||||
assert not leave_task_box["task"].done()
|
||||
return result
|
||||
|
||||
db.scalar = scalar_then_race # type: ignore[method-assign]
|
||||
yield db
|
||||
|
||||
monkeypatch.setattr(scope_backfill_mod, "tracked_db", tracked_db_with_leave_race)
|
||||
|
||||
ok = await scope_backfill_mod._copy_chunk( # pyright: ignore[reportPrivateUsage]
|
||||
workspace_name,
|
||||
scope_peer.name,
|
||||
session.name,
|
||||
[
|
||||
scope_backfill_mod._CopySpec( # pyright: ignore[reportPrivateUsage]
|
||||
observed=sender.name,
|
||||
source_id=source.id,
|
||||
content=source.content,
|
||||
embedding=None,
|
||||
internal_metadata={},
|
||||
times_derived=1,
|
||||
source_ids=None,
|
||||
session_name=session.name,
|
||||
)
|
||||
],
|
||||
store_in_postgres=True,
|
||||
)
|
||||
assert ok is True
|
||||
|
||||
leave_task = leave_task_box["task"]
|
||||
await asyncio.wait_for(leave_task, timeout=2.0)
|
||||
assert leave_finished.is_set()
|
||||
|
||||
copies = await _get_docs(
|
||||
db_session,
|
||||
workspace_name,
|
||||
observer=scope_peer.name,
|
||||
observed=sender.name,
|
||||
include_deleted=False,
|
||||
)
|
||||
assert len(copies) == 1
|
||||
assert copies[0].internal_metadata.get(COPIED_FROM_KEY) == source.id
|
||||
|
||||
membership = await db_session.scalar(
|
||||
select(models.SessionPeer.left_at).where(
|
||||
models.SessionPeer.workspace_name == workspace_name,
|
||||
models.SessionPeer.session_name == session.name,
|
||||
models.SessionPeer.peer_name == scope_peer.name,
|
||||
)
|
||||
)
|
||||
assert membership is not None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 3. Multi-peer session
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -956,3 +1082,87 @@ async def test_backfill_status_writes_preserve_the_scope_kind_flag(
|
|||
await db_session.commit()
|
||||
metadata = await assert_still_a_scope("clearing the status")
|
||||
assert session_name not in metadata.get("backfill_status", {})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_backfill_embeds_and_writes_in_bounded_chunks(
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""Phases 2-4 run per chunk, so a large session never holds every vector."""
|
||||
from src.deriver import scope_backfill
|
||||
from src.embedding_client import embedding_client
|
||||
|
||||
test_workspace, sender = sample_data
|
||||
workspace_name = test_workspace.name
|
||||
scope_name = str(generate_nanoid())
|
||||
scope_peer = await _create_scope_peer(db_session, workspace_name, scope_name)
|
||||
session = await _create_session(db_session, workspace_name)
|
||||
await _join_scope(db_session, workspace_name, session.name, scope_peer.name)
|
||||
await _create_collection(
|
||||
db_session, workspace_name, observer=sender.name, observed=sender.name
|
||||
)
|
||||
await _create_collection(
|
||||
db_session, workspace_name, observer=scope_peer.name, observed=sender.name
|
||||
)
|
||||
for i in range(3):
|
||||
source = await _create_document(
|
||||
db_session,
|
||||
workspace_name,
|
||||
observer=sender.name,
|
||||
observed=sender.name,
|
||||
session_name=session.name,
|
||||
content=f"fact {i}",
|
||||
)
|
||||
source.embedding = None
|
||||
await db_session.commit()
|
||||
|
||||
batch_sizes: list[int] = []
|
||||
seen_specs: list[scope_backfill._CopySpec] = [] # pyright: ignore[reportPrivateUsage]
|
||||
peak_live_embeddings = 0
|
||||
original_embed = embedding_client.simple_batch_embed
|
||||
original_copy_chunk = scope_backfill._copy_chunk # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
async def recording_embed(texts: list[str], **kwargs: Any) -> list[list[float]]:
|
||||
batch_sizes.append(len(texts))
|
||||
return await original_embed(texts, **kwargs)
|
||||
|
||||
async def counting_copy_chunk(
|
||||
ws_name: str,
|
||||
peer_name: str,
|
||||
sess_name: str,
|
||||
plans: list[scope_backfill._CopySpec], # pyright: ignore[reportPrivateUsage]
|
||||
store_in_postgres: bool,
|
||||
) -> bool:
|
||||
nonlocal peak_live_embeddings
|
||||
seen_specs.extend(plans)
|
||||
result = await original_copy_chunk(
|
||||
ws_name, peer_name, sess_name, plans, store_in_postgres
|
||||
)
|
||||
# Sampled after this chunk syncs but before _run_backfill drops its
|
||||
# vectors, so every *earlier* chunk must already be cleared and the
|
||||
# live count can never exceed one chunk. That drop is the whole
|
||||
# memory bound; without it this peaks at 3 instead of 2.
|
||||
peak_live_embeddings = max(
|
||||
peak_live_embeddings,
|
||||
sum(1 for spec in seen_specs if spec.embedding is not None),
|
||||
)
|
||||
return result
|
||||
|
||||
monkeypatch.setattr(scope_backfill, "BACKFILL_CHUNK_SIZE", 2)
|
||||
monkeypatch.setattr(embedding_client, "simple_batch_embed", recording_embed)
|
||||
monkeypatch.setattr(scope_backfill, "_copy_chunk", counting_copy_chunk)
|
||||
|
||||
await process_scope_backfill(
|
||||
ScopeBackfillPayload(scope_peer=scope_peer.name, session_name=session.name),
|
||||
workspace_name,
|
||||
)
|
||||
|
||||
assert batch_sizes == [2, 1]
|
||||
assert peak_live_embeddings == 2
|
||||
copies = await _get_docs(
|
||||
db_session, workspace_name, observer=scope_peer.name, observed=sender.name
|
||||
)
|
||||
assert len(copies) == 3
|
||||
assert all(copy.embedding is not None for copy in copies)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,321 @@
|
|||
"""Tests for the read-only count of collections whose next dream is due."""
|
||||
|
||||
import datetime
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from nanoid import generate as generate_nanoid
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from src import models
|
||||
from src.dreamer.dream_due import count_due_dreams
|
||||
from src.schemas import DreamType
|
||||
from src.utils.work_unit import construct_work_unit_key
|
||||
|
||||
|
||||
def _now() -> datetime.datetime:
|
||||
return datetime.datetime.now(datetime.UTC)
|
||||
|
||||
|
||||
async def _make_collection(
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
internal_metadata: dict[str, object] | None = None,
|
||||
) -> models.Collection:
|
||||
workspace, peer = sample_data
|
||||
collection = models.Collection(
|
||||
observer=peer.name,
|
||||
observed=peer.name,
|
||||
workspace_name=workspace.name,
|
||||
internal_metadata=internal_metadata or {},
|
||||
)
|
||||
db_session.add(collection)
|
||||
await db_session.commit()
|
||||
return collection
|
||||
|
||||
|
||||
async def _make_session(
|
||||
db_session: AsyncSession,
|
||||
workspace_name: str,
|
||||
configuration: dict[str, object] | None = None,
|
||||
) -> str:
|
||||
session = models.Session(
|
||||
name=f"s-{generate_nanoid()}",
|
||||
workspace_name=workspace_name,
|
||||
configuration=configuration or {},
|
||||
)
|
||||
db_session.add(session)
|
||||
await db_session.commit()
|
||||
return session.name
|
||||
|
||||
|
||||
async def _insert_docs(
|
||||
db_session: AsyncSession,
|
||||
collection: models.Collection,
|
||||
level: str,
|
||||
count: int,
|
||||
*,
|
||||
age_minutes: int = 0,
|
||||
session_name: str | None = None,
|
||||
sessionless: bool = False,
|
||||
) -> None:
|
||||
if session_name is None and not sessionless:
|
||||
session_name = await _make_session(db_session, collection.workspace_name)
|
||||
created_at = _now() - datetime.timedelta(minutes=age_minutes)
|
||||
for _ in range(count):
|
||||
db_session.add(
|
||||
models.Document(
|
||||
content="test",
|
||||
level=level,
|
||||
workspace_name=collection.workspace_name,
|
||||
observer=collection.observer,
|
||||
observed=collection.observed,
|
||||
session_name=session_name,
|
||||
created_at=created_at,
|
||||
)
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
|
||||
async def _insert_dream_item(
|
||||
db_session: AsyncSession,
|
||||
collection: models.Collection,
|
||||
*,
|
||||
age_minutes: int,
|
||||
processed: bool,
|
||||
error: str | None = None,
|
||||
) -> None:
|
||||
work_unit_key = construct_work_unit_key(
|
||||
collection.workspace_name,
|
||||
{
|
||||
"task_type": "dream",
|
||||
"observer": collection.observer,
|
||||
"observed": collection.observed,
|
||||
"dream_type": DreamType.OMNI.value,
|
||||
},
|
||||
)
|
||||
db_session.add(
|
||||
models.QueueItem(
|
||||
work_unit_key=work_unit_key,
|
||||
payload={"task_type": "dream"},
|
||||
task_type="dream",
|
||||
workspace_name=collection.workspace_name,
|
||||
processed=processed,
|
||||
error=error,
|
||||
created_at=_now() - datetime.timedelta(minutes=age_minutes),
|
||||
)
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _pin_dream_config(): # pyright: ignore[reportUnusedFunction]
|
||||
with (
|
||||
patch("src.dreamer.dream_due.settings.DREAM.ENABLED", True),
|
||||
patch("src.dreamer.dream_due.settings.DREAM.DOCUMENT_THRESHOLD", 50),
|
||||
patch("src.dreamer.dream_due.settings.DREAM.ENABLED_TYPES", ["omni"]),
|
||||
patch("src.dreamer.dream_due.settings.DREAM.IDLE_TIMEOUT_MINUTES", 60),
|
||||
patch("src.dreamer.dream_due.settings.DREAM.MIN_HOURS_BETWEEN_DREAMS", 8),
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestCountDueDreams:
|
||||
async def test_below_threshold_is_not_due(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
collection = await _make_collection(db_session, sample_data)
|
||||
await _insert_docs(db_session, collection, "explicit", 30, age_minutes=90)
|
||||
|
||||
assert await count_due_dreams(db_session) == 0
|
||||
|
||||
async def test_derived_levels_do_not_count(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
collection = await _make_collection(db_session, sample_data)
|
||||
await _insert_docs(db_session, collection, "explicit", 30, age_minutes=90)
|
||||
await _insert_docs(db_session, collection, "deductive", 40, age_minutes=90)
|
||||
|
||||
assert await count_due_dreams(db_session) == 0
|
||||
|
||||
async def test_threshold_met_but_not_idle_is_not_due(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
"""A collection still receiving documents is not idle yet."""
|
||||
collection = await _make_collection(db_session, sample_data)
|
||||
await _insert_docs(db_session, collection, "explicit", 60, age_minutes=1)
|
||||
|
||||
assert await count_due_dreams(db_session) == 0
|
||||
|
||||
async def test_threshold_met_and_idle_is_due(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
collection = await _make_collection(db_session, sample_data)
|
||||
await _insert_docs(db_session, collection, "explicit", 60, age_minutes=90)
|
||||
|
||||
assert await count_due_dreams(db_session) == 1
|
||||
|
||||
async def test_documents_since_last_dream_uses_stored_count(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
collection = await _make_collection(
|
||||
db_session, sample_data, {"dream": {"last_dream_document_count": 40}}
|
||||
)
|
||||
await _insert_docs(db_session, collection, "explicit", 60, age_minutes=90)
|
||||
|
||||
assert await count_due_dreams(db_session) == 0
|
||||
|
||||
async def test_min_hours_gate_blocks_a_recent_dream(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
last_dream_at = (_now() - datetime.timedelta(hours=2)).isoformat()
|
||||
collection = await _make_collection(
|
||||
db_session, sample_data, {"dream": {"last_dream_at": last_dream_at}}
|
||||
)
|
||||
await _insert_docs(db_session, collection, "explicit", 60, age_minutes=90)
|
||||
|
||||
assert await count_due_dreams(db_session) == 0
|
||||
|
||||
async def test_naive_last_dream_at_is_read_as_utc(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
"""A stored timestamp with no offset must gate, not raise."""
|
||||
naive = (_now() - datetime.timedelta(hours=2)).replace(tzinfo=None).isoformat()
|
||||
collection = await _make_collection(
|
||||
db_session, sample_data, {"dream": {"last_dream_at": naive}}
|
||||
)
|
||||
await _insert_docs(db_session, collection, "explicit", 60, age_minutes=90)
|
||||
|
||||
assert await count_due_dreams(db_session) == 0
|
||||
|
||||
async def test_pending_dream_item_blocks(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
collection = await _make_collection(db_session, sample_data)
|
||||
await _insert_docs(db_session, collection, "explicit", 60, age_minutes=90)
|
||||
await _insert_dream_item(
|
||||
db_session, collection, age_minutes=10, processed=False
|
||||
)
|
||||
|
||||
assert await count_due_dreams(db_session) == 0
|
||||
|
||||
async def test_failed_dream_waits_for_new_documents(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
"""Without this the count never returns to zero."""
|
||||
collection = await _make_collection(db_session, sample_data)
|
||||
await _insert_docs(db_session, collection, "explicit", 60, age_minutes=90)
|
||||
await _insert_dream_item(
|
||||
db_session, collection, age_minutes=80, processed=True, error="boom"
|
||||
)
|
||||
|
||||
assert await count_due_dreams(db_session) == 0
|
||||
|
||||
async def test_failed_dream_retries_after_new_documents(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
collection = await _make_collection(db_session, sample_data)
|
||||
await _insert_docs(db_session, collection, "explicit", 60, age_minutes=90)
|
||||
await _insert_dream_item(
|
||||
db_session, collection, age_minutes=80, processed=True, error="boom"
|
||||
)
|
||||
await _insert_docs(db_session, collection, "explicit", 1, age_minutes=70)
|
||||
|
||||
assert await count_due_dreams(db_session) == 1
|
||||
|
||||
async def test_sessionless_documents_are_not_due(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
"""The deriver's own enqueue path refuses these, so they must not count."""
|
||||
collection = await _make_collection(db_session, sample_data)
|
||||
await _insert_docs(
|
||||
db_session, collection, "explicit", 60, age_minutes=90, sessionless=True
|
||||
)
|
||||
|
||||
assert await count_due_dreams(db_session) == 0
|
||||
|
||||
async def test_newest_document_decides_the_session(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
collection = await _make_collection(db_session, sample_data)
|
||||
await _insert_docs(db_session, collection, "explicit", 60, age_minutes=120)
|
||||
|
||||
assert await count_due_dreams(db_session) == 1
|
||||
|
||||
await _insert_docs(
|
||||
db_session, collection, "explicit", 1, age_minutes=90, sessionless=True
|
||||
)
|
||||
|
||||
assert await count_due_dreams(db_session) == 0
|
||||
|
||||
async def test_session_with_dreams_disabled_is_not_due(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
"""A dream the enqueue path would refuse must not be counted."""
|
||||
collection = await _make_collection(db_session, sample_data)
|
||||
session_name = await _make_session(
|
||||
db_session,
|
||||
collection.workspace_name,
|
||||
{"dream": {"enabled": False}},
|
||||
)
|
||||
await _insert_docs(
|
||||
db_session,
|
||||
collection,
|
||||
"explicit",
|
||||
60,
|
||||
age_minutes=90,
|
||||
session_name=session_name,
|
||||
)
|
||||
|
||||
assert await count_due_dreams(db_session) == 0
|
||||
|
||||
async def test_dreams_disabled_globally_returns_zero(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
collection = await _make_collection(db_session, sample_data)
|
||||
await _insert_docs(db_session, collection, "explicit", 60, age_minutes=90)
|
||||
|
||||
with patch("src.dreamer.dream_due.settings.DREAM.ENABLED", False):
|
||||
assert await count_due_dreams(db_session) == 0
|
||||
|
||||
async def test_card_refresh_is_never_counted(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
sample_data: tuple[models.Workspace, models.Peer],
|
||||
):
|
||||
collection = await _make_collection(db_session, sample_data)
|
||||
await _insert_docs(db_session, collection, "explicit", 60, age_minutes=90)
|
||||
|
||||
with patch(
|
||||
"src.dreamer.dream_due.settings.DREAM.ENABLED_TYPES", ["card_refresh"]
|
||||
):
|
||||
assert await count_due_dreams(db_session) == 0
|
||||
|
|
@ -277,6 +277,40 @@ async def test_peer_chat_non_streaming(
|
|||
assert response is None or isinstance(response, str)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("timeout", [None, 2.5])
|
||||
async def test_peer_chat_forwards_per_call_timeout(
|
||||
client_fixture: tuple[Honcho, str],
|
||||
timeout: float | None,
|
||||
) -> None:
|
||||
honcho_client, client_type = client_fixture
|
||||
timeout_label = "default" if timeout is None else "override"
|
||||
|
||||
if client_type == "async":
|
||||
peer = await honcho_client.aio.peer(id=f"test-timeout-{timeout_label}-async")
|
||||
|
||||
async def mock_post(*args: object, **kwargs: object) -> dict[str, str]: # pyright: ignore[reportUnusedParameter]
|
||||
return {"content": "ok"}
|
||||
|
||||
with patch.object(
|
||||
peer._honcho._async_http_client, # pyright: ignore[reportPrivateUsage]
|
||||
"post",
|
||||
side_effect=mock_post,
|
||||
) as mock:
|
||||
result = await peer.aio.chat("What do I like?", timeout=timeout)
|
||||
else:
|
||||
peer = honcho_client.peer(id=f"test-timeout-{timeout_label}-sync")
|
||||
with patch.object(
|
||||
peer._honcho._http, # pyright: ignore[reportPrivateUsage]
|
||||
"post",
|
||||
return_value={"content": "ok"},
|
||||
) as mock:
|
||||
result = peer.chat("What do I like?", timeout=timeout)
|
||||
|
||||
assert result == "ok"
|
||||
assert mock.call_args.kwargs["timeout"] == timeout
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_peer_representation_no_params(
|
||||
client_fixture: tuple[Honcho, str],
|
||||
|
|
|
|||
|
|
@ -131,6 +131,19 @@ def test_deriver_token_combos_are_valid_and_complete():
|
|||
) not in ingestion
|
||||
|
||||
|
||||
_API_DERIVER_METRIC_GAUGES = (
|
||||
"deriver_outstanding_work_seconds",
|
||||
"deriver_queue_work_units_eligible",
|
||||
"deriver_queue_work_units_claimed",
|
||||
"deriver_queue_items_pending",
|
||||
"deriver_queue_oldest_pending_age_seconds",
|
||||
"dreams_due",
|
||||
"message_embeddings_pending_due",
|
||||
)
|
||||
|
||||
_SHARED_DERIVER_METRIC_GAUGES = ("message_embeddings_pending",)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# API-process zero-init
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -161,6 +174,8 @@ def test_api_init_materializes_dialectic_and_embed():
|
|||
)
|
||||
assert sample("embed_now_tasks_shed_total") is not None
|
||||
assert sample("embed_now_tasks_in_flight") == 0.0 # gauge, explicit .set(0)
|
||||
for gauge in (*_API_DERIVER_METRIC_GAUGES, *_SHARED_DERIVER_METRIC_GAUGES):
|
||||
assert sample(gauge) == 0.0, f"{gauge} was not zero-initialized"
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("metrics_enabled")
|
||||
|
|
@ -310,6 +325,9 @@ def test_deriver_init_does_not_touch_api_counters():
|
|||
# the API-process embed_now counters are equally off-limits
|
||||
assert sample("embed_now_tasks_shed_total") is None
|
||||
assert sample("embed_now_tasks_in_flight") is None
|
||||
# so are the deriver-work gauges: the deriver never measures its own backlog
|
||||
for gauge in _API_DERIVER_METRIC_GAUGES:
|
||||
assert sample(gauge) is None, f"{gauge} must be API-only"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -0,0 +1,207 @@
|
|||
"""Tests for the outstanding-work value, the poller and the JSON route."""
|
||||
|
||||
import time
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
from src import schemas
|
||||
from src.backlog import (
|
||||
DeriverMetricsPoller,
|
||||
DeriverMetricsSnapshot,
|
||||
active_work_seconds,
|
||||
outstanding_work_seconds,
|
||||
)
|
||||
from src.routers import deriver_metrics
|
||||
|
||||
|
||||
class TestScaleSignal:
|
||||
def test_nothing_outstanding_reads_zero(self):
|
||||
assert outstanding_work_seconds(schemas.DeriverMetrics(), dreams_due=0) == 0.0
|
||||
|
||||
def test_claimable_work_reports_the_active_value(self):
|
||||
stats = schemas.DeriverMetrics(eligible_work_units=1)
|
||||
|
||||
assert outstanding_work_seconds(stats, dreams_due=0) == active_work_seconds()
|
||||
|
||||
def test_work_in_flight_still_reports_the_active_value(self):
|
||||
"""A row claimed a moment ago has a small age and would read as idle."""
|
||||
stats = schemas.DeriverMetrics(
|
||||
claimed_work_units=1, pending_items=1, oldest_pending_age_seconds=2.0
|
||||
)
|
||||
|
||||
assert outstanding_work_seconds(stats, dreams_due=0) == active_work_seconds()
|
||||
|
||||
def test_waiting_batch_reports_its_real_age(self):
|
||||
"""The real age is what tells a caller how close the flush is."""
|
||||
stats = schemas.DeriverMetrics(
|
||||
pending_items=3, oldest_pending_age_seconds=1234.0
|
||||
)
|
||||
|
||||
assert outstanding_work_seconds(stats, dreams_due=0) == 1234.0
|
||||
|
||||
def test_embeddings_due_an_attempt_report_the_active_value(self):
|
||||
stats = schemas.DeriverMetrics(embeddings_pending=5, embeddings_pending_due=5)
|
||||
|
||||
assert outstanding_work_seconds(stats, dreams_due=0) == active_work_seconds()
|
||||
|
||||
def test_embeddings_inside_their_retry_wait_do_not(self):
|
||||
"""Otherwise one permanently failing row holds the value up for hours."""
|
||||
stats = schemas.DeriverMetrics(embeddings_pending=5)
|
||||
|
||||
assert outstanding_work_seconds(stats, dreams_due=0) == 0.0
|
||||
|
||||
def test_a_due_dream_reports_the_active_value(self):
|
||||
assert (
|
||||
outstanding_work_seconds(schemas.DeriverMetrics(), dreams_due=1)
|
||||
== active_work_seconds()
|
||||
)
|
||||
|
||||
def test_active_value_is_positive(self):
|
||||
assert active_work_seconds() > 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestPoller:
|
||||
async def test_refresh_publishes_a_snapshot(self):
|
||||
stats = schemas.DeriverMetrics(eligible_work_units=2, pending_items=4)
|
||||
poller = DeriverMetricsPoller()
|
||||
|
||||
with (
|
||||
patch(
|
||||
"src.backlog.crud.get_deriver_metrics",
|
||||
AsyncMock(return_value=stats),
|
||||
),
|
||||
patch("src.backlog.count_due_dreams", AsyncMock(return_value=3)),
|
||||
):
|
||||
await poller.refresh()
|
||||
|
||||
snapshot = poller.snapshot
|
||||
assert snapshot.measured_at is not None
|
||||
assert snapshot.stats.eligible_work_units == 2
|
||||
assert snapshot.dreams_due == 3
|
||||
assert snapshot.signal_seconds == active_work_seconds()
|
||||
|
||||
async def test_dream_query_runs_on_its_own_spacing(self):
|
||||
"""The dream query is the expensive one, so it must not run every pass."""
|
||||
stats = schemas.DeriverMetrics()
|
||||
poller = DeriverMetricsPoller()
|
||||
dream_count = AsyncMock(return_value=1)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"src.backlog.crud.get_deriver_metrics",
|
||||
AsyncMock(return_value=stats),
|
||||
),
|
||||
patch("src.backlog.count_due_dreams", dream_count),
|
||||
):
|
||||
await poller.refresh()
|
||||
await poller.refresh()
|
||||
|
||||
assert dream_count.await_count == 1
|
||||
assert poller.snapshot.dreams_due == 1
|
||||
|
||||
async def test_a_failed_dream_query_is_retried_on_the_next_pass(self):
|
||||
"""Advancing the deadline first would republish the old count for a whole interval."""
|
||||
stats = schemas.DeriverMetrics()
|
||||
poller = DeriverMetricsPoller()
|
||||
dream_count = AsyncMock(side_effect=[RuntimeError("db down"), 4])
|
||||
|
||||
with (
|
||||
patch(
|
||||
"src.backlog.crud.get_deriver_metrics",
|
||||
AsyncMock(return_value=stats),
|
||||
),
|
||||
patch("src.backlog.count_due_dreams", dream_count),
|
||||
):
|
||||
with pytest.raises(RuntimeError):
|
||||
await poller.refresh()
|
||||
await poller.refresh()
|
||||
|
||||
assert dream_count.await_count == 2
|
||||
assert poller.snapshot.dreams_due == 4
|
||||
|
||||
async def test_a_failed_pass_leaves_the_previous_snapshot_alone(self):
|
||||
"""A half-finished pass must never be published as a measurement."""
|
||||
stats = schemas.DeriverMetrics(eligible_work_units=1)
|
||||
poller = DeriverMetricsPoller()
|
||||
|
||||
with (
|
||||
patch(
|
||||
"src.backlog.crud.get_deriver_metrics",
|
||||
AsyncMock(return_value=stats),
|
||||
),
|
||||
patch("src.backlog.count_due_dreams", AsyncMock(return_value=0)),
|
||||
):
|
||||
await poller.refresh()
|
||||
|
||||
first = poller.snapshot
|
||||
|
||||
with (
|
||||
patch(
|
||||
"src.backlog.crud.get_deriver_metrics",
|
||||
AsyncMock(side_effect=RuntimeError("db down")),
|
||||
),
|
||||
pytest.raises(RuntimeError),
|
||||
):
|
||||
await poller.refresh()
|
||||
|
||||
assert poller.snapshot is first
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestDeriverMetricsRoute:
|
||||
async def test_serves_the_cached_snapshot(self):
|
||||
poller = DeriverMetricsPoller()
|
||||
poller._snapshot = DeriverMetricsSnapshot( # pyright: ignore[reportPrivateUsage]
|
||||
signal_seconds=1800.0,
|
||||
dreams_due=1,
|
||||
stats=schemas.DeriverMetrics(eligible_work_units=2, pending_items=5),
|
||||
measured_at=time.time(),
|
||||
)
|
||||
deriver_metrics.set_deriver_metrics_poller(poller)
|
||||
try:
|
||||
body = await deriver_metrics.get_deriver_metrics_response()
|
||||
finally:
|
||||
deriver_metrics.set_deriver_metrics_poller(None)
|
||||
|
||||
assert body["outstanding_work_seconds"] == 1800.0
|
||||
assert body["eligible_work_units"] == 2
|
||||
assert body["pending_items"] == 5
|
||||
assert body["dreams_due"] == 1
|
||||
|
||||
async def test_errors_before_the_first_pass(self):
|
||||
"""A 503 tells the caller there is no measurement; a 0 would be a lie."""
|
||||
deriver_metrics.set_deriver_metrics_poller(DeriverMetricsPoller())
|
||||
try:
|
||||
with pytest.raises(HTTPException) as excinfo:
|
||||
await deriver_metrics.get_deriver_metrics_response()
|
||||
finally:
|
||||
deriver_metrics.set_deriver_metrics_poller(None)
|
||||
|
||||
assert excinfo.value.status_code == 503
|
||||
|
||||
async def test_serves_an_old_snapshot_with_its_age(self):
|
||||
"""The caller decides what is too old, from measurement_age_seconds."""
|
||||
poller = DeriverMetricsPoller()
|
||||
poller._snapshot = DeriverMetricsSnapshot( # pyright: ignore[reportPrivateUsage]
|
||||
signal_seconds=7.0,
|
||||
measured_at=time.time() - 3600,
|
||||
)
|
||||
deriver_metrics.set_deriver_metrics_poller(poller)
|
||||
try:
|
||||
body = await deriver_metrics.get_deriver_metrics_response()
|
||||
finally:
|
||||
deriver_metrics.set_deriver_metrics_poller(None)
|
||||
|
||||
assert body["outstanding_work_seconds"] == 7.0
|
||||
assert body["measurement_age_seconds"] >= 3600
|
||||
|
||||
async def test_errors_when_no_poller_is_registered(self):
|
||||
deriver_metrics.set_deriver_metrics_poller(None)
|
||||
|
||||
with pytest.raises(HTTPException) as excinfo:
|
||||
await deriver_metrics.get_deriver_metrics_response()
|
||||
|
||||
assert excinfo.value.status_code == 503
|
||||
|
|
@ -5,9 +5,11 @@ from pydantic import ValidationError
|
|||
|
||||
from src.config import settings
|
||||
from src.schemas import (
|
||||
DialecticOptions,
|
||||
DocumentCreate,
|
||||
DocumentMetadata,
|
||||
MessageCreate,
|
||||
ObservationInput,
|
||||
PeerCreate,
|
||||
ReasoningConfiguration,
|
||||
ResolvedConfiguration,
|
||||
|
|
@ -275,3 +277,68 @@ class TestReasoningCustomInstructionsValidation:
|
|||
configuration = ReasoningConfiguration(custom_instructions=custom_instructions)
|
||||
|
||||
assert configuration.custom_instructions == custom_instructions
|
||||
|
||||
|
||||
class TestNulByteSanitization:
|
||||
"""Postgres rejects NUL (0x00) in text columns and in jsonb strings.
|
||||
|
||||
Models emit these as `\\u0000` escapes in tool-call arguments, which the
|
||||
JSON parser decodes into real NUL bytes, so model-generated text needs the
|
||||
same treatment as user-supplied input.
|
||||
"""
|
||||
|
||||
def test_document_content_strips_nul(self):
|
||||
document = DocumentCreate(
|
||||
content="the key is at c:\\\x00users\\amal",
|
||||
metadata=DocumentMetadata(message_ids=[1], message_created_at="2026-08-28"),
|
||||
embedding=[0.1],
|
||||
)
|
||||
|
||||
assert document.content == "the key is at c:\\users\\amal"
|
||||
|
||||
def test_all_nul_document_content_is_rejected_not_emptied(self):
|
||||
"""The validator runs before `min_length`, so content that is nothing
|
||||
but NUL fails validation rather than being stored as an empty string."""
|
||||
with pytest.raises(ValidationError):
|
||||
DocumentCreate(
|
||||
content="\x00\x00",
|
||||
metadata=DocumentMetadata(
|
||||
message_ids=[1], message_created_at="2026-08-28"
|
||||
),
|
||||
embedding=[0.1],
|
||||
)
|
||||
|
||||
def test_message_content_strips_nul(self):
|
||||
message = MessageCreate(peer_id="peer", content="before\x00after")
|
||||
|
||||
assert message.content == "beforeafter"
|
||||
|
||||
def test_metadata_strips_nul_at_every_depth(self):
|
||||
message = MessageCreate(
|
||||
peer_id="peer",
|
||||
content="hi",
|
||||
metadata={"a\x00b": {"c": ["d\x00e", 1]}},
|
||||
)
|
||||
|
||||
assert message.metadata == {"ab": {"c": ["de", 1]}}
|
||||
|
||||
def test_observation_content_strips_nul(self):
|
||||
observation = ObservationInput(content="before\x00after")
|
||||
|
||||
assert observation.content == "beforeafter"
|
||||
|
||||
def test_all_nul_observation_content_is_rejected_not_emptied(self):
|
||||
"""Sanitization runs before `min_length`, so an all-NUL observation is
|
||||
reported back to the model as a validation failure rather than saved
|
||||
as an empty document."""
|
||||
with pytest.raises(ValidationError):
|
||||
ObservationInput(content="\x00\x00")
|
||||
|
||||
def test_all_nul_query_is_rejected_not_emptied(self):
|
||||
"""`NulStripped` runs before the field's own constraints, so a query
|
||||
that is nothing but NUL fails `min_length` instead of reaching the
|
||||
dialectic as an empty prompt."""
|
||||
options = DialecticOptions.model_validate({"query": "before\x00after"})
|
||||
assert options.query == "beforeafter"
|
||||
with pytest.raises(ValidationError):
|
||||
DialecticOptions.model_validate({"query": "\x00"})
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ import asyncio
|
|||
import json
|
||||
from collections.abc import Callable
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
|
|
@ -82,7 +82,7 @@ async def workspace_test_data(
|
|||
await db_session.flush()
|
||||
|
||||
# Create messages
|
||||
now = datetime.now(timezone.utc)
|
||||
now = datetime.now(UTC)
|
||||
messages: list[models.Message] = []
|
||||
for i in range(6):
|
||||
peer_name = [peer1.name, peer2.name, peer3.name][i % 3]
|
||||
|
|
@ -593,7 +593,7 @@ class TestSearchMemoryWorkspace:
|
|||
content="I really like programming in Python",
|
||||
seq_in_session=1,
|
||||
token_count=10,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
created_at=datetime.now(UTC),
|
||||
)
|
||||
db_session.add(msg)
|
||||
await db_session.flush()
|
||||
|
|
@ -919,7 +919,7 @@ class TestGetObservationContextWorkspace:
|
|||
content="LEAKED_FROM_OTHER_SESSION",
|
||||
seq_in_session=messages[0].seq_in_session,
|
||||
token_count=10,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
created_at=datetime.now(UTC),
|
||||
)
|
||||
db_session.add(leaked_message)
|
||||
await db_session.commit()
|
||||
|
|
@ -1253,3 +1253,63 @@ class TestWorkspaceChatPrompt:
|
|||
}
|
||||
assert agent.messages[0]["content"] == workspace_agent_system_prompt(offered)
|
||||
assert agent._prefetch_heading() == "Workspace overview (prefetched)" # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
def test_forbids_clarifying_questions(self) -> None:
|
||||
"""The endpoint is non-interactive, so the prompt must say so.
|
||||
|
||||
Without this the model answers a recall query with a plan and a menu of
|
||||
lookups for a caller that cannot reply. The pair agent talks to a peer
|
||||
and is deliberately left alone.
|
||||
"""
|
||||
from src.dialectic.prompts import (
|
||||
agent_system_prompt,
|
||||
workspace_agent_system_prompt,
|
||||
)
|
||||
|
||||
prompt = workspace_agent_system_prompt()
|
||||
assert "NO CLARIFYING QUESTIONS" in prompt
|
||||
assert "NO CLARIFYING QUESTIONS" not in agent_system_prompt(
|
||||
"alice", "alice", None, None
|
||||
)
|
||||
|
||||
|
||||
class TestWorkspaceToolChoice:
|
||||
"""The workspace agent must search before it answers.
|
||||
|
||||
Its prefetch is an orientation overview, not the corpus, so a turn with no
|
||||
tool call ends the loop with whatever the overview happened to contain.
|
||||
"""
|
||||
|
||||
@pytest.mark.parametrize("level", ["minimal", "low", "medium", "high", "max"])
|
||||
def test_first_turn_requires_a_tool_call(self, level: str) -> None:
|
||||
from src.config import settings
|
||||
from src.dialectic.workspace import WorkspaceDialecticAgent
|
||||
|
||||
agent = WorkspaceDialecticAgent(workspace_name="w", reasoning_level=level) # pyright: ignore[reportArgumentType]
|
||||
level_settings = settings.DIALECTIC.LEVELS[level] # pyright: ignore[reportArgumentType]
|
||||
assert agent._tool_choice(level_settings) == "required" # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
def test_pair_agent_keeps_the_configured_choice(self) -> None:
|
||||
from src.config import settings
|
||||
from src.dialectic.core import DialecticAgent
|
||||
|
||||
agent = DialecticAgent(
|
||||
workspace_name="w", session_name=None, observer="a", observed="a"
|
||||
)
|
||||
level_settings = settings.DIALECTIC.LEVELS["low"]
|
||||
assert (
|
||||
agent._tool_choice(level_settings) # pyright: ignore[reportPrivateUsage]
|
||||
== level_settings.TOOL_CHOICE
|
||||
)
|
||||
|
||||
def test_a_configured_non_auto_choice_is_passed_through(self) -> None:
|
||||
from src.config import DialecticLevelSettings, settings
|
||||
from src.dialectic.workspace import WorkspaceDialecticAgent
|
||||
|
||||
agent = WorkspaceDialecticAgent(workspace_name="w")
|
||||
pinned = DialecticLevelSettings(
|
||||
MODEL_CONFIG=settings.DIALECTIC.LEVELS["low"].MODEL_CONFIG,
|
||||
MAX_TOOL_ITERATIONS=5,
|
||||
TOOL_CHOICE="none",
|
||||
)
|
||||
assert agent._tool_choice(pinned) == "none" # pyright: ignore[reportPrivateUsage]
|
||||
|
|
|
|||
|
|
@ -0,0 +1,117 @@
|
|||
"""DB-free unit tests for src/utils/retryable_errors.py."""
|
||||
|
||||
import asyncio
|
||||
from typing import cast
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from sqlalchemy.exc import DBAPIError, OperationalError
|
||||
|
||||
from src.utils.retryable_errors import is_retryable_db_error, is_retryable_error
|
||||
|
||||
|
||||
class FakePGError(Exception):
|
||||
"""Stands in for a driver exception carrying a SQLSTATE."""
|
||||
|
||||
sqlstate: str | None
|
||||
|
||||
def __init__(self, sqlstate: str | None) -> None:
|
||||
super().__init__(f"fake pg error ({sqlstate})")
|
||||
self.sqlstate = sqlstate
|
||||
|
||||
|
||||
def _dbapi_error(
|
||||
sqlstate: str | None,
|
||||
*,
|
||||
orig: BaseException | None = None,
|
||||
connection_invalidated: bool = False,
|
||||
) -> DBAPIError:
|
||||
if orig is None and sqlstate is not None:
|
||||
orig = FakePGError(sqlstate)
|
||||
return OperationalError(
|
||||
"SELECT 1",
|
||||
{},
|
||||
cast(BaseException, orig),
|
||||
connection_invalidated=connection_invalidated,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("sqlstate", "expected"),
|
||||
[
|
||||
("40P01", True), # deadlock_detected
|
||||
("40001", True), # serialization_failure
|
||||
("55P03", True), # lock_not_available
|
||||
("57014", True), # query_canceled
|
||||
("08006", True), # connection_failure
|
||||
("23505", False), # unique_violation
|
||||
("42P01", False), # undefined_table
|
||||
("22P02", False), # invalid_text_representation
|
||||
],
|
||||
)
|
||||
def test_sqlstate_classification(sqlstate: str, expected: bool):
|
||||
exc = _dbapi_error(sqlstate)
|
||||
assert is_retryable_db_error(exc) is expected
|
||||
assert is_retryable_error(exc) is expected
|
||||
|
||||
|
||||
def test_orig_none_is_terminal():
|
||||
assert not is_retryable_db_error(_dbapi_error(None))
|
||||
|
||||
|
||||
def test_sqlstate_on_orig_cause():
|
||||
"""SQLSTATE found by walking orig.__cause__ when orig itself has none."""
|
||||
wrapper = Exception("driver wrapper")
|
||||
wrapper.__cause__ = FakePGError("40P01")
|
||||
assert is_retryable_db_error(_dbapi_error(None, orig=wrapper))
|
||||
|
||||
|
||||
def test_connection_invalidated_is_retryable():
|
||||
exc = _dbapi_error(None, connection_invalidated=True)
|
||||
assert is_retryable_db_error(exc)
|
||||
|
||||
|
||||
def test_dbapi_error_nested_in_cause_chain():
|
||||
outer = RuntimeError("save failed")
|
||||
outer.__cause__ = _dbapi_error("40P01")
|
||||
assert is_retryable_db_error(outer)
|
||||
assert is_retryable_error(outer)
|
||||
|
||||
|
||||
def test_non_db_exceptions_are_not_db_retryable():
|
||||
assert not is_retryable_db_error(ValueError("bad input"))
|
||||
assert not is_retryable_db_error(httpx.ConnectTimeout("timed out"))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("exc", "expected"),
|
||||
[
|
||||
(httpx.ConnectTimeout("timed out"), True),
|
||||
(httpx.ReadTimeout("timed out"), True),
|
||||
(httpx.ConnectError("connection refused"), True),
|
||||
(ConnectionResetError("reset"), True),
|
||||
(asyncio.TimeoutError(), True),
|
||||
(TimeoutError(), True),
|
||||
(ValueError("bad input"), False),
|
||||
(httpx.HTTPStatusError("401", request=None, response=None), False), # pyright: ignore[reportArgumentType]
|
||||
],
|
||||
)
|
||||
def test_transport_classification(exc: BaseException, expected: bool):
|
||||
assert is_retryable_error(exc) is expected
|
||||
assert not is_retryable_db_error(exc)
|
||||
|
||||
|
||||
def test_transport_error_nested_in_cause_chain():
|
||||
"""SDK wrappers (e.g. APIConnectionError) chain to httpx via __cause__."""
|
||||
wrapper = RuntimeError("provider call failed")
|
||||
wrapper.__cause__ = httpx.ConnectError("connection refused")
|
||||
assert is_retryable_error(wrapper)
|
||||
assert not is_retryable_db_error(wrapper)
|
||||
|
||||
|
||||
def test_cause_cycle_terminates():
|
||||
a = RuntimeError("a")
|
||||
b = RuntimeError("b")
|
||||
a.__cause__ = b
|
||||
b.__cause__ = a
|
||||
assert not is_retryable_error(a)
|
||||
|
|
@ -0,0 +1,38 @@
|
|||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from src.utils.sanitization import strip_nul
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("value", "expected"),
|
||||
[
|
||||
pytest.param("before\x00after", "beforeafter", id="string"),
|
||||
pytest.param("no nul here", "no nul here", id="string-unchanged"),
|
||||
pytest.param("\x00\x00", "", id="string-all-nul"),
|
||||
pytest.param(["a\x00b", "c"], ["ab", "c"], id="list"),
|
||||
pytest.param({"k\x00": "v\x00"}, {"k": "v"}, id="dict-key-and-value"),
|
||||
pytest.param(
|
||||
{"a": [{"b": "c\x00d"}]},
|
||||
{"a": [{"b": "cd"}]},
|
||||
id="nested",
|
||||
),
|
||||
# Optional fields are passed in without a guard, so None has to survive.
|
||||
pytest.param(None, None, id="none"),
|
||||
pytest.param(7, 7, id="int"),
|
||||
pytest.param(True, True, id="bool"),
|
||||
pytest.param([], [], id="empty-list"),
|
||||
],
|
||||
)
|
||||
def test_strip_nul(value: Any, expected: Any) -> None:
|
||||
assert strip_nul(value) == expected
|
||||
|
||||
|
||||
def test_strip_nul_does_not_mutate_its_argument() -> None:
|
||||
original = {"a": ["b\x00c"]}
|
||||
|
||||
stripped = strip_nul(original)
|
||||
|
||||
assert stripped == {"a": ["bc"]}
|
||||
assert original == {"a": ["b\x00c"]}
|
||||
Loading…
Reference in New Issue