From b7bcb327385894748224a53fe44c3cb3dea28845 Mon Sep 17 00:00:00 2001 From: Erosika Date: Tue, 25 Aug 2026 16:44:37 -0400 Subject: [PATCH 01/17] feat(docs): load the GTM container on every docs page DEV-2465 step 1. Mintlify injects gtm.js on all docs pages; the container is audited to be inert on /docs before this merges, so the snippet loads and nothing fires. Cookiebot and GA4 arrive later as container publishes, consent first. Merging this publishes the docs within minutes, so it stays unmerged until Marc confirms the container audit. --- docs/docs.json | 74 ++++++++++++++++++++++++++++++++++++++------------ 1 file changed, 57 insertions(+), 17 deletions(-) diff --git a/docs/docs.json b/docs/docs.json index b9d0e498..e0873acb 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -19,14 +19,21 @@ }, "favicon": "/favicon.svg", "contextual": { - "options": ["copy", "view", "chatgpt", "claude"] + "options": [ + "copy", + "view", + "chatgpt", + "claude" + ] }, "navigation": { "versions": [ { "version": "v3.1.0", "api": { - "openapi": ["v3/openapi.json"] + "openapi": [ + "v3/openapi.json" + ] }, "tabs": [ { @@ -90,7 +97,9 @@ "groups": [ { "group": "Overview", - "pages": ["v3/guides/overview"] + "pages": [ + "v3/guides/overview" + ] }, { "group": "Integrations", @@ -130,7 +139,9 @@ }, { "group": "Migrations", - "pages": ["v3/guides/migrations/mem0"] + "pages": [ + "v3/guides/migrations/mem0" + ] } ] }, @@ -160,7 +171,9 @@ "groups": [ { "group": "API Documentation", - "pages": ["v3/api-reference/introduction"] + "pages": [ + "v3/api-reference/introduction" + ] }, { "group": "workspaces", @@ -238,7 +251,9 @@ }, { "group": "miscellaneous", - "pages": ["v3/api-reference/endpoint/keys/create-key"] + "pages": [ + "v3/api-reference/endpoint/keys/create-key" + ] } ] }, @@ -259,7 +274,9 @@ { "version": "v2.5.1", "api": { - "openapi": ["v2/openapi.json"] + "openapi": [ + "v2/openapi.json" + ] }, "tabs": [ { @@ -306,11 +323,15 @@ "groups": [ { "group": "Getting Started", - "pages": ["v2/guides/overview"] + "pages": [ + "v2/guides/overview" + ] }, { "group": "Migrations", - "pages": ["v2/migrations/from-mem0"] + "pages": [ + "v2/migrations/from-mem0" + ] }, { "group": "Integrations", @@ -335,7 +356,9 @@ "groups": [ { "group": "API Documentation", - "pages": ["v2/api-reference/introduction"] + "pages": [ + "v2/api-reference/introduction" + ] }, { "group": "workspaces", @@ -439,7 +462,9 @@ { "version": "v1.1.0", "api": { - "openapi": ["openapi.json"] + "openapi": [ + "openapi.json" + ] }, "tabs": [ { @@ -469,15 +494,23 @@ "groups": [ { "group": "Getting Started", - "pages": ["v1/guides/overview", "v1/guides/streaming-response"] + "pages": [ + "v1/guides/overview", + "v1/guides/streaming-response" + ] }, { "group": "Application Interfaces", - "pages": ["v1/guides/discord", "v1/guides/honcho-mcp"] + "pages": [ + "v1/guides/discord", + "v1/guides/honcho-mcp" + ] }, { "group": "Personal Memory", - "pages": ["v1/guides/dialectic-endpoint"] + "pages": [ + "v1/guides/dialectic-endpoint" + ] } ] }, @@ -486,7 +519,9 @@ "groups": [ { "group": "API Documentation", - "pages": ["v1/api-reference/introduction"] + "pages": [ + "v1/api-reference/introduction" + ] }, { "group": "apps", @@ -534,7 +569,9 @@ }, { "group": "keys", - "pages": ["v1/api-reference/endpoint/keys/create-key"] + "pages": [ + "v1/api-reference/endpoint/keys/create-key" + ] }, { "group": "metamessages", @@ -595,6 +632,9 @@ "integrations": { "posthog": { "apiKey": "phc_1yrzzcgywqXGcerkkI4g7C0YfyPMcAKNOOvGcjTCiUk" + }, + "gtm": { + "tagId": "GTM-NSPT9PJF" } } -} +} \ No newline at end of file From 0b9ae0017009af27e662356407d8ed9c565cbb4c Mon Sep 17 00:00:00 2001 From: Erosika Date: Tue, 25 Aug 2026 16:45:18 -0400 Subject: [PATCH 02/17] feat(docs): PostHog loads only with a granting consent answer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DEV-2465 open question 5, option d. Mintlify's built-in integration loaded PostHog unconditionally on all 249 docs pages — a visitor who declined on the homepage was tracked one click later in the docs. The integration key comes out of docs.json; docs/posthog-consent.js loads PostHog directly instead, only when the CookieConsent cookie grants Statistics (or holds Cookiebot's -1 marker), and listens for the consent events so a grant on the docs banner itself loads it too. Trade recorded on the ticket: this bypasses the ph.mintlify.com proxy, so ad blockers reduce docs PostHog volume. Verify after deploy that Mintlify's page CSP allows us-assets.i.posthog.com; if it blocks, fall back to option c. --- docs/docs.json | 3 --- docs/posthog-consent.js | 43 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 43 insertions(+), 3 deletions(-) create mode 100644 docs/posthog-consent.js diff --git a/docs/docs.json b/docs/docs.json index e0873acb..b1dad2ba 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -630,9 +630,6 @@ } }, "integrations": { - "posthog": { - "apiKey": "phc_1yrzzcgywqXGcerkkI4g7C0YfyPMcAKNOOvGcjTCiUk" - }, "gtm": { "tagId": "GTM-NSPT9PJF" } diff --git a/docs/posthog-consent.js b/docs/posthog-consent.js new file mode 100644 index 00000000..214c6b1c --- /dev/null +++ b/docs/posthog-consent.js @@ -0,0 +1,43 @@ +// 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(/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.onload = function () { + 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', + }) + } + document.head.appendChild(s) + } + + if (granted()) { + loadPosthog() + return + } + // A grant made on the docs banner itself (step 2) loads it live. + var events = ['CookiebotOnConsentReady', 'CookiebotOnAccept'] + for (var i = 0; i < events.length; i++) { + window.addEventListener(events[i], function () { + if (granted()) loadPosthog() + }) + } +})() From b5d1a1ae540774e78a48cd66269c9a2cb5719c03 Mon Sep 17 00:00:00 2001 From: Erosika Date: Wed, 26 Aug 2026 10:25:26 -0400 Subject: [PATCH 03/17] fix(docs): loader survives a failed fetch and honors withdrawal MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review findings on #1071. The cookie match now requires the exact CookieConsent name boundary. A failed array.js request resets the loaded flag so later consent events retry. And consent events now run a full sync: withdrawal opts an already running instance out, and a re-grant opts it back in — same behavior as the landing site's gate. --- docs/posthog-consent.js | 38 +++++++++++++++++++++++++++++--------- 1 file changed, 29 insertions(+), 9 deletions(-) diff --git a/docs/posthog-consent.js b/docs/posthog-consent.js index 214c6b1c..3e93da82 100644 --- a/docs/posthog-consent.js +++ b/docs/posthog-consent.js @@ -5,7 +5,7 @@ var loaded = false function granted() { - var m = document.cookie.match(/CookieConsent=([^;]*)/) + var m = document.cookie.match(/(?:^|;\s*)CookieConsent=([^;]*)/) if (!m) return false var v = decodeURIComponent(m[1]) // "-1" is Cookiebot's consent-not-required marker. @@ -18,6 +18,9 @@ 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 () { window.posthog.init(KEY, { api_host: 'https://us.i.posthog.com', @@ -29,15 +32,32 @@ document.head.appendChild(s) } - if (granted()) { - loadPosthog() - return + 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() + } } - // A grant made on the docs banner itself (step 2) loads it live. - var events = ['CookiebotOnConsentReady', 'CookiebotOnAccept'] + + sync() + var events = [ + 'CookiebotOnConsentReady', + 'CookiebotOnAccept', + 'CookiebotOnDecline', + ] for (var i = 0; i < events.length; i++) { - window.addEventListener(events[i], function () { - if (granted()) loadPosthog() - }) + window.addEventListener(events[i], sync) } })() From f3db11ef3a05671a6e1c4779cb2a85e105152d56 Mon Sep 17 00:00:00 2001 From: Erosika Date: Thu, 27 Aug 2026 10:53:10 -0400 Subject: [PATCH 04/17] chore(docs): undo array reformatting in docs.json The integrations change is the only intended edit. The one-item arrays go back to their single-line form and the trailing newline returns. --- docs/docs.json | 71 ++++++++++++-------------------------------------- 1 file changed, 17 insertions(+), 54 deletions(-) diff --git a/docs/docs.json b/docs/docs.json index b1dad2ba..12b5fe01 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -19,21 +19,14 @@ }, "favicon": "/favicon.svg", "contextual": { - "options": [ - "copy", - "view", - "chatgpt", - "claude" - ] + "options": ["copy", "view", "chatgpt", "claude"] }, "navigation": { "versions": [ { "version": "v3.1.0", "api": { - "openapi": [ - "v3/openapi.json" - ] + "openapi": ["v3/openapi.json"] }, "tabs": [ { @@ -97,9 +90,7 @@ "groups": [ { "group": "Overview", - "pages": [ - "v3/guides/overview" - ] + "pages": ["v3/guides/overview"] }, { "group": "Integrations", @@ -139,9 +130,7 @@ }, { "group": "Migrations", - "pages": [ - "v3/guides/migrations/mem0" - ] + "pages": ["v3/guides/migrations/mem0"] } ] }, @@ -171,9 +160,7 @@ "groups": [ { "group": "API Documentation", - "pages": [ - "v3/api-reference/introduction" - ] + "pages": ["v3/api-reference/introduction"] }, { "group": "workspaces", @@ -251,9 +238,7 @@ }, { "group": "miscellaneous", - "pages": [ - "v3/api-reference/endpoint/keys/create-key" - ] + "pages": ["v3/api-reference/endpoint/keys/create-key"] } ] }, @@ -274,9 +259,7 @@ { "version": "v2.5.1", "api": { - "openapi": [ - "v2/openapi.json" - ] + "openapi": ["v2/openapi.json"] }, "tabs": [ { @@ -323,15 +306,11 @@ "groups": [ { "group": "Getting Started", - "pages": [ - "v2/guides/overview" - ] + "pages": ["v2/guides/overview"] }, { "group": "Migrations", - "pages": [ - "v2/migrations/from-mem0" - ] + "pages": ["v2/migrations/from-mem0"] }, { "group": "Integrations", @@ -356,9 +335,7 @@ "groups": [ { "group": "API Documentation", - "pages": [ - "v2/api-reference/introduction" - ] + "pages": ["v2/api-reference/introduction"] }, { "group": "workspaces", @@ -462,9 +439,7 @@ { "version": "v1.1.0", "api": { - "openapi": [ - "openapi.json" - ] + "openapi": ["openapi.json"] }, "tabs": [ { @@ -494,23 +469,15 @@ "groups": [ { "group": "Getting Started", - "pages": [ - "v1/guides/overview", - "v1/guides/streaming-response" - ] + "pages": ["v1/guides/overview", "v1/guides/streaming-response"] }, { "group": "Application Interfaces", - "pages": [ - "v1/guides/discord", - "v1/guides/honcho-mcp" - ] + "pages": ["v1/guides/discord", "v1/guides/honcho-mcp"] }, { "group": "Personal Memory", - "pages": [ - "v1/guides/dialectic-endpoint" - ] + "pages": ["v1/guides/dialectic-endpoint"] } ] }, @@ -519,9 +486,7 @@ "groups": [ { "group": "API Documentation", - "pages": [ - "v1/api-reference/introduction" - ] + "pages": ["v1/api-reference/introduction"] }, { "group": "apps", @@ -569,9 +534,7 @@ }, { "group": "keys", - "pages": [ - "v1/api-reference/endpoint/keys/create-key" - ] + "pages": ["v1/api-reference/endpoint/keys/create-key"] }, { "group": "metamessages", @@ -634,4 +597,4 @@ "tagId": "GTM-NSPT9PJF" } } -} \ No newline at end of file +} From 86f8eb3e6ddae9c79592c794f574377b474e3b96 Mon Sep 17 00:00:00 2001 From: Erosika Date: Thu, 27 Aug 2026 10:53:10 -0400 Subject: [PATCH 05/17] fix(docs): loader re-checks consent before init and captures SPA pageviews If consent is withdrawn while array.js downloads, sync() runs before window.posthog exists and the opt-out is skipped. onload now re-checks granted() and resets loaded so a later re-grant retries. Mintlify swaps pages without a reload, so capture_pageview: 'history_change' records navigation past the landing page. --- docs/posthog-consent.js | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docs/posthog-consent.js b/docs/posthog-consent.js index 3e93da82..acb47860 100644 --- a/docs/posthog-consent.js +++ b/docs/posthog-consent.js @@ -22,11 +22,17 @@ 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) From 03253d7a088df185c5a441463d5333b60cd4d49c Mon Sep 17 00:00:00 2001 From: Eugene Eisenstein Date: Mon, 31 Aug 2026 13:32:43 -0400 Subject: [PATCH 06/17] fix(deriver): strip NUL bytes from model-generated observations (#1095) * fix(deriver): strip NUL bytes from model-generated observations Postgres rejects NUL (0x00) in text columns and in jsonb strings. API ingress has always stripped it from user-supplied content, but the deriver's own output did not go through any equivalent: a model can emit a \u0000 escape in its tool-call arguments, which the JSON parser decodes into a real NUL byte. Seen in production when models transcribe shell output (`tr '\x00' '\n'`) or Windows paths (`c:\users\amal`). The NUL reached the exact-content dedup pre-fetch in create_documents as a bind parameter, so the query raised DataError before any row was written and the whole batch for that observer was dropped. Strip in _normalized_observation and _normalized_observation_input -- the points that already normalize text for persistence and embedding -- so the embedded text matches the stored text. premises and sources are covered too, since they ride along in internal_metadata. The emptiness check now runs after normalization, because str.strip() does not remove NUL and all-NUL content would otherwise be stored as an empty string. DocumentCreate.content gets a mode="before" validator as a backstop for callers that bypass those paths; running before the length constraint makes all-NUL content fail min_length rather than silently empty out. The NUL helpers move out of schemas/api.py into utils/sanitization.py as a single recursive strip_nul, so ingress and internal paths share one implementation. It is overloaded to keep str -> str for the callers that chain .strip(), and passes None through so optional fields need no guard. Fixes HONCHO-4XZ * fix: broaden nul strip check * chore: code simplification --------- Co-authored-by: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com> --- src/crud/representation.py | 29 ++++- src/schemas/api.py | 31 +---- src/schemas/internal.py | 12 +- src/utils/agent_tools.py | 25 +++- src/utils/sanitization.py | 46 +++++++ tests/crud/test_representation_manager.py | 145 +++++++++++++++++++--- tests/test_schema_validations.py | 67 ++++++++++ tests/utils/test_sanitization.py | 38 ++++++ 8 files changed, 334 insertions(+), 59 deletions(-) create mode 100644 src/utils/sanitization.py create mode 100644 tests/utils/test_sanitization.py diff --git a/src/crud/representation.py b/src/crud/representation.py index 6fafb842..85b3e931 100644 --- a/src/crud/representation.py +++ b/src/crud/representation.py @@ -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") diff --git a/src/schemas/api.py b/src/schemas/api.py index 43d91b26..34258d99 100644 --- a/src/schemas/api.py +++ b/src/schemas/api.py @@ -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: diff --git a/src/schemas/internal.py b/src/schemas/internal.py index e014431f..2d299feb 100644 --- a/src/schemas/internal.py +++ b/src/schemas/internal.py @@ -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.""" diff --git a/src/utils/agent_tools.py b/src/utils/agent_tools.py index b753c462..de07e38f 100644 --- a/src/utils/agent_tools.py +++ b/src/utils/agent_tools.py @@ -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") diff --git a/src/utils/sanitization.py b/src/utils/sanitization.py new file mode 100644 index 00000000..4879cc72 --- /dev/null +++ b/src/utils/sanitization.py @@ -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) diff --git a/tests/crud/test_representation_manager.py b/tests/crud/test_representation_manager.py index 3f3d6f40..9392b78e 100644 --- a/tests/crud/test_representation_manager.py +++ b/tests/crud/test_representation_manager.py @@ -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(), ) diff --git a/tests/test_schema_validations.py b/tests/test_schema_validations.py index 97b8b64a..e3a285f2 100644 --- a/tests/test_schema_validations.py +++ b/tests/test_schema_validations.py @@ -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"}) diff --git a/tests/utils/test_sanitization.py b/tests/utils/test_sanitization.py new file mode 100644 index 00000000..dc84647f --- /dev/null +++ b/tests/utils/test_sanitization.py @@ -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"]} From c300236c110c6e544ced07c843c5daf44fa133a5 Mon Sep 17 00:00:00 2001 From: Aakash Kattelu Date: Tue, 1 Sep 2026 09:17:34 -0400 Subject: [PATCH 07/17] fix(deriver): reduce scope backfill memory usage (#1104) * fix(deriver): chunk scope backfill so large sessions don't OOM the worker _run_backfill embedded, wrote, and synced every planned copy at once, holding one Python float list per document. A 14k-document session is ~580MB of vectors alone, and several backfills run concurrently, which OOM-killed the deriver at its 1000Mi limit and crash-looped it since the work units never completed. Phases 2-4 now run per chunk of 500 specs and drop each chunk's embeddings once synced. Co-Authored-By: Claude Fable 5 * fix(deriver): hydrate backfill embeddings per chunk Phase 1 no longer materializes every source embedding into plans. load_only skips the vector column on the plan queries, and each chunk reloads only its source embeddings before embed/write/sync. * fix(deriver): lock scope membership across backfill chunk writes SELECT ... FOR UPDATE on the active SessionPeer row so a concurrent leave cannot commit between the membership check and the copy inserts. Adds a concurrency test that asserts the leave blocks until commit. * fix: add test for memory bound --------- Co-authored-by: Claude Fable 5 Co-authored-by: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com> --- src/deriver/scope_backfill.py | 114 ++++++++++++-- tests/deriver/test_scope_backfill.py | 212 ++++++++++++++++++++++++++- 2 files changed, 310 insertions(+), 16 deletions(-) diff --git a/src/deriver/scope_backfill.py b/src/deriver/scope_backfill.py index 81017e28..9be082f9 100644 --- a/src/deriver/scope_backfill.py +++ b/src/deriver/scope_backfill.py @@ -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( diff --git a/tests/deriver/test_scope_backfill.py b/tests/deriver/test_scope_backfill.py index 06464c6f..efca34f3 100644 --- a/tests/deriver/test_scope_backfill.py +++ b/tests/deriver/test_scope_backfill.py @@ -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) From a026bebdef91e2b0d052574a653afc39b3ad3918 Mon Sep 17 00:00:00 2001 From: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com> Date: Tue, 1 Sep 2026 23:03:44 -0400 Subject: [PATCH 08/17] chore(docs): Add explanation on deleting data and cloud vs local differences (#1114) --- docs/docs.json | 3 +- .../endpoint/keys/create-key.mdx | 10 ++ .../features/advanced/deleting-data.mdx | 131 ++++++++++++++++++ .../features/advanced/overview.mdx | 1 + .../features/advanced/webhooks.mdx | 8 ++ docs/v3/documentation/reference/platform.mdx | 2 + docs/v3/documentation/reference/sdk.mdx | 6 + 7 files changed, 160 insertions(+), 1 deletion(-) create mode 100644 docs/v3/documentation/features/advanced/deleting-data.mdx diff --git a/docs/docs.json b/docs/docs.json index de5fa522..f130a939 100644 --- a/docs/docs.json +++ b/docs/docs.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" ] } ] diff --git a/docs/v3/api-reference/endpoint/keys/create-key.mdx b/docs/v3/api-reference/endpoint/keys/create-key.mdx index 484229c9..9f9b0470 100644 --- a/docs/v3/api-reference/endpoint/keys/create-key.mdx +++ b/docs/v3/api-reference/endpoint/keys/create-key.mdx @@ -1,3 +1,13 @@ --- openapi: post /v3/keys --- + + +**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. + diff --git a/docs/v3/documentation/features/advanced/deleting-data.mdx b/docs/v3/documentation/features/advanced/deleting-data.mdx new file mode 100644 index 00000000..1a6ab444 --- /dev/null +++ b/docs/v3/documentation/features/advanced/deleting-data.mdx @@ -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. + + +```python Python +session.delete() +``` + +```typescript TypeScript +await session.delete(); +``` + + +## 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. + + +```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"); +``` + + +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: + + +```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); +} +``` + + +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. diff --git a/docs/v3/documentation/features/advanced/overview.mdx b/docs/v3/documentation/features/advanced/overview.mdx index d0d7cc2c..8c160734 100644 --- a/docs/v3/documentation/features/advanced/overview.mdx +++ b/docs/v3/documentation/features/advanced/overview.mdx @@ -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 diff --git a/docs/v3/documentation/features/advanced/webhooks.mdx b/docs/v3/documentation/features/advanced/webhooks.mdx index 8976750a..ae7d90ad 100644 --- a/docs/v3/documentation/features/advanced/webhooks.mdx +++ b/docs/v3/documentation/features/advanced/webhooks.mdx @@ -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. + +**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. + + ## Registering an Endpoint diff --git a/docs/v3/documentation/reference/platform.mdx b/docs/v3/documentation/reference/platform.mdx index 3ed501d5..28042d4c 100644 --- a/docs/v3/documentation/reference/platform.mdx +++ b/docs/v3/documentation/reference/platform.mdx @@ -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. diff --git a/docs/v3/documentation/reference/sdk.mdx b/docs/v3/documentation/reference/sdk.mdx index 77652bf0..432a4aab 100644 --- a/docs/v3/documentation/reference/sdk.mdx +++ b/docs/v3/documentation/reference/sdk.mdx @@ -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); ``` From ccdb8ba11341752e61265c5fc3fe49d3312791cf Mon Sep 17 00:00:00 2001 From: Aakash Kattelu Date: Wed, 2 Sep 2026 11:07:50 -0400 Subject: [PATCH 09/17] fix(deriver): fix create_documents deadlock (#1033) * fix(deriver): eliminate create_documents deadlock and stop silently burning batches on transient errors Two concurrent work units writing the same (workspace, observer, observed) collection deadlocked on times_derived reinforcement UPDATEs issued in batch order (DEV-1975, 682 events in 90 days). The deadlock was swallowed per-document, the loop cascaded PendingRollbackErrors against the dead session, the whole batch was lost, and the queue item was marked processed. - serialize writers per collection with a transaction-scoped advisory lock (pg_advisory_xact_lock + SET LOCAL lock_timeout), skipped for insert-only batches; covers all three row-lock sites in one move - hoist external-vector-store dup-candidate resolution ahead of the first DB statement so the lock's critical section contains no network calls - abort the batch on SQLAlchemyError instead of continuing through an aborted transaction; per-document skip semantics kept for non-DB errors - classify transient errors (new src/utils/retryable_errors.py) and retry them via a bounded in-process counter instead of marking items errored * fix(deriver): replace create_documents advisory lock with id-ordered row locks Advisory locks are database-scoped and would serialize every writer to a collection, including across Groudon tenants that share names. Collect reinforcement and replace ops during the loop, lock target rows with SELECT ... ORDER BY id FOR UPDATE, then apply. populate_existing reloads times_derived so a prefetched identity-map row cannot lose a concurrent increment. * fix(deriver): harden create_documents candidate hoist and test isolation Skip empty embeddings on the external-store path, isolate per-document resolve failures, and keep replacement times_derived in the in-batch ledger. Patch get_external_vector_store in the hoist test and cover in-loop SQLAlchemyError abort. * fix(deriver): address CodeRabbit findings on create_documents deadlock fix - Distinguish external resolve failure ([] skip) from pgvector fallback (None) so _semantic_dup_decision never re-enters external I/O under an open session - Bound external candidate hoist concurrency with a semaphore - Map in-loop IntegrityError to ValidationException for a uniform contract - Persist transient retry attempts on the oldest unprocessed queue item so every deriver instance shares one MAX_RETRYABLE_ATTEMPTS budget - Cover resolve-failure skip and multi-manager reclaim of the retry budget * fix(deriver): harden retry metadata cleanup and stale reinforce fallback - Strip _retry_attempts from payloads in the same transaction as mark_queue_items_as_processed / mark_queue_item_as_errored - Clear shared retry metadata only after a successful terminal mark - On reinforce, if the locked target is gone or soft-deleted, insert the incoming document instead of dropping it - Skip pgvector semantic lookup when embedding is empty so query_documents cannot embed under an open session * fix(deriver): address review on deadlock retry and row-lock apply Strip _retry_attempts before payload validation so non-representation tasks are not burned as extra_forbidden. Re-raise retryable observer save errors after telemetry so the queue actually retries. Skip same-batch reinforce fallbacks after a replace. Revert unordered FOR UPDATE on mark processed/errored and drop post-commit retry cleanup from the success path. * fix: add test and simplify queue query --------- Co-authored-by: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com> --- CLAUDE.md | 10 + src/crud/document.py | 456 +++++++++++----- src/deriver/consumer.py | 7 +- src/deriver/deriver.py | 7 + src/deriver/queue_manager.py | 160 +++++- src/utils/queue_payload.py | 8 + src/utils/retryable_errors.py | 86 +++ tests/crud/test_document.py | 642 ++++++++++++++++++++++- tests/deriver/test_deriver_processing.py | 71 ++- tests/deriver/test_queue_processing.py | 367 ++++++++++++- tests/utils/test_retryable_errors.py | 117 +++++ 11 files changed, 1765 insertions(+), 166 deletions(-) create mode 100644 src/utils/retryable_errors.py create mode 100644 tests/utils/test_retryable_errors.py diff --git a/CLAUDE.md b/CLAUDE.md index 6a83066d..ab842f99 100644 --- a/CLAUDE.md +++ b/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 `.** 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 ` 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. diff --git a/src/crud/document.py b/src/crud/document.py index 37eb94b4..7a3fc63a 100644 --- a/src/crud/document.py +++ b/src/crud/document.py @@ -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 ) diff --git a/src/deriver/consumer.py b/src/deriver/consumer.py index 118135ca..6251dc73 100644 --- a/src/deriver/consumer.py +++ b/src/deriver/consumer.py @@ -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 diff --git a/src/deriver/deriver.py b/src/deriver/deriver.py index f76c4d52..4c1e4dd2 100644 --- a/src/deriver/deriver.py +++ b/src/deriver/deriver.py @@ -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}" diff --git a/src/deriver/queue_manager.py b/src/deriver/queue_manager.py index 1c19c131..b98c0ef6 100644 --- a/src/deriver/queue_manager.py +++ b/src/deriver/queue_manager.py @@ -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, or_, select, update from sqlalchemy.dialects.postgresql import insert from sqlalchemy.engine import CursorResult from sqlalchemy.ext.asyncio import AsyncSession @@ -43,6 +43,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 +55,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 +309,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 ) @@ -591,11 +599,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 +624,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 +667,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 +733,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 +763,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 +1125,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: diff --git a/src/utils/queue_payload.py b/src/utils/queue_payload.py index 59815ca5..6f3f1105 100644 --- a/src/utils/queue_payload.py +++ b/src/utils/queue_payload.py @@ -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.""" diff --git a/src/utils/retryable_errors.py b/src/utils/retryable_errors.py new file mode 100644 index 00000000..357f594d --- /dev/null +++ b/src/utils/retryable_errors.py @@ -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) + ) diff --git a/tests/crud/test_document.py b/tests/crud/test_document.py index 6686e688..593ee52b 100644 --- a/tests/crud/test_document.py +++ b/tests/crud/test_document.py @@ -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() diff --git a/tests/deriver/test_deriver_processing.py b/tests/deriver/test_deriver_processing.py index 29a983b5..86ab9d33 100644 --- a/tests/deriver/test_deriver_processing.py +++ b/tests/deriver/test_deriver_processing.py @@ -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 diff --git a/tests/deriver/test_queue_processing.py b/tests/deriver/test_queue_processing.py index 81dd0632..ce8e3e13 100644 --- a/tests/deriver/test_queue_processing.py +++ b/tests/deriver/test_queue_processing.py @@ -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 diff --git a/tests/utils/test_retryable_errors.py b/tests/utils/test_retryable_errors.py new file mode 100644 index 00000000..bead2d2e --- /dev/null +++ b/tests/utils/test_retryable_errors.py @@ -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) From ced151420001e6c237eb148dd939c362b29617ed Mon Sep 17 00:00:00 2001 From: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:19:42 -0400 Subject: [PATCH 10/17] =?UTF-8?q?chore(docs):=20Add=20section=20about=20ha?= =?UTF-8?q?rness=20integrations=20and=20deepseek=20harn=E2=80=A6=20(#1116)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * chore(docs): Add section about harness integrations and deepseek harness to docs * chore: Add section about harness integrations and deepseek harness to docs --- README.md | 49 ++++- docs/docs.json | 1 + .../guides/integrations/deepseek-harness.mdx | 198 ++++++++++++++++++ 3 files changed, 245 insertions(+), 3 deletions(-) create mode 100644 docs/v3/guides/integrations/deepseek-harness.mdx diff --git a/README.md b/README.md index 2d354c66..056253cd 100644 --- a/README.md +++ b/README.md @@ -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 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 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 diff --git a/docs/docs.json b/docs/docs.json index f130a939..83b2475a 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -104,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", diff --git a/docs/v3/guides/integrations/deepseek-harness.mdx b/docs/v3/guides/integrations/deepseek-harness.mdx new file mode 100644 index 00000000..fa612f6c --- /dev/null +++ b/docs/v3/guides/integrations/deepseek-harness.mdx @@ -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 + + +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. + + +```bash +dsh plugin --profile 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. + + +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. + + +## 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 - + "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 + } + } + } +} +``` + + +Unsupported or renamed keys are reported at startup rather than silently ignored, so a stale config tells you what it is no longer doing. + + +### 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) | `-` | Stable across restarts and branches | +| `per-repo` | `-` | Same memory from any subdirectory | +| `git-branch` | `--` | Falls back to `per-directory` outside a repo or on a detached HEAD | +| `per-session` | `-chat-` | A clean slate every restart | +| `global` | `` | One memory for everything | + + +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. + + +### 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 + + + + Source code, issues, and README. + + + + Learn about peers, sessions, and dialectic reasoning. + + From 997b4764b926b13c23b72a034d5c96c1ac364862 Mon Sep 17 00:00:00 2001 From: Aakash Kattelu Date: Wed, 2 Sep 2026 12:37:17 -0400 Subject: [PATCH 11/17] chore: add changelog and version updates (#1117) API: 3.1.0 -> 3.1.1 Python/TS SDK: 2.4.0 -> 2.4.0 (unchanged) CLI: 0.1.4 -> 0.1.4 (unchanged) --- CHANGELOG.md | 15 +++++++++++++++ README.md | 2 +- docs/changelog/compatibility-guide.mdx | 3 ++- docs/changelog/introduction.mdx | 23 +++++++++++++++++++---- docs/docs.json | 2 +- docs/v3/openapi.json | 2 +- pyproject.toml | 2 +- uv.lock | 2 +- 8 files changed, 41 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3bef42b6..96a8534f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,21 @@ 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/). +## [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 diff --git a/README.md b/README.md index 056253cd..e9011cef 100644 --- a/README.md +++ b/README.md @@ -8,7 +8,7 @@ --- -![Static Badge](https://img.shields.io/badge/Server-3.1.0-blue) +![Static Badge](https://img.shields.io/badge/Server-3.1.1-blue) [![PyPI version](https://img.shields.io/pypi/v/honcho-ai.svg)](https://pypi.org/project/honcho-ai/) [![NPM version](https://img.shields.io/npm/v/@honcho-ai/sdk.svg)](https://npmjs.org/package/@honcho-ai/sdk) [![CLI](https://img.shields.io/pypi/v/honcho-cli.svg?label=honcho-cli)](https://pypi.org/project/honcho-cli/) diff --git a/docs/changelog/compatibility-guide.mdx b/docs/changelog/compatibility-guide.mdx index abdd1361..20cf904f 100644 --- a/docs/changelog/compatibility-guide.mdx +++ b/docs/changelog/compatibility-guide.mdx @@ -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 | diff --git a/docs/changelog/introduction.mdx b/docs/changelog/introduction.mdx index a152d16a..264f683b 100644 --- a/docs/changelog/introduction.mdx +++ b/docs/changelog/introduction.mdx @@ -27,7 +27,22 @@ Welcome to the Honcho changelog! This section documents all notable changes to t ### Honcho API and SDK Changelogs - + + ### 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) + + + ### 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 [Python SDK](https://pypi.org/project/honcho-ai/) - + ### 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 [TypeScript SDK](https://www.npmjs.com/package/@honcho-ai/sdk) - + ### 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 [Honcho CLI](https://pypi.org/project/honcho-cli/) - + ### 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` diff --git a/docs/docs.json b/docs/docs.json index 83b2475a..a1dbb487 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -28,7 +28,7 @@ "navigation": { "versions": [ { - "version": "v3.1.0", + "version": "v3.1.1", "api": { "openapi": ["v3/openapi.json"] }, diff --git a/docs/v3/openapi.json b/docs/v3/openapi.json index b2bc0adf..e7a836ee 100644 --- a/docs/v3/openapi.json +++ b/docs/v3/openapi.json @@ -9,7 +9,7 @@ "url": "https://honcho.dev/", "email": "hello@plasticlabs.ai" }, - "version": "3.1.0" + "version": "3.1.1" }, "servers": [ { diff --git a/pyproject.toml b/pyproject.toml index a6681ceb..7c858e23 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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"}, diff --git a/uv.lock b/uv.lock index 0e84663f..a9a92a6d 100644 --- a/uv.lock +++ b/uv.lock @@ -824,7 +824,7 @@ wheels = [ [[package]] name = "honcho" -version = "3.1.0" +version = "3.1.1" source = { virtual = "." } dependencies = [ { name = "alembic" }, From 5d992bc65afcfbc05a5911ab4edbaa88ef64c690 Mon Sep 17 00:00:00 2001 From: Ulysse Pence <736903+ulyssepence@users.noreply.github.com> Date: Wed, 2 Sep 2026 13:42:48 -0400 Subject: [PATCH 12/17] feat(api): Export deriver backlog as metrics from API endpoint (#1115) --- src/backlog.py | 146 +++++++++ src/config.py | 3 + src/crud/__init__.py | 7 +- src/crud/deriver.py | 152 ++++++++- src/deriver/queue_manager.py | 45 +-- src/dreamer/dream_due.py | 216 +++++++++++++ src/main.py | 12 + src/reconciler/embed_now.py | 4 +- src/reconciler/sync_vectors.py | 8 +- src/routers/deriver_metrics.py | 41 +++ src/schemas/__init__.py | 2 + src/schemas/internal.py | 11 + src/telemetry/prometheus/metrics.py | 107 ++++++ tests/crud/test_deriver_metrics_query.py | 394 +++++++++++++++++++++++ tests/dreamer/test_dream_due.py | 321 ++++++++++++++++++ tests/telemetry/test_metric_zero_init.py | 18 ++ tests/test_deriver_metrics.py | 207 ++++++++++++ 17 files changed, 1656 insertions(+), 38 deletions(-) create mode 100644 src/backlog.py create mode 100644 src/dreamer/dream_due.py create mode 100644 src/routers/deriver_metrics.py create mode 100644 tests/crud/test_deriver_metrics_query.py create mode 100644 tests/dreamer/test_dream_due.py create mode 100644 tests/test_deriver_metrics.py diff --git a/src/backlog.py b/src/backlog.py new file mode 100644 index 00000000..1180be57 --- /dev/null +++ b/src/backlog.py @@ -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 + ) diff --git a/src/config.py b/src/config.py index 993f9bfb..80827327 100644 --- a/src/config.py +++ b/src/config.py @@ -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 diff --git a/src/crud/__init__.py b/src/crud/__init__.py index 0e920717..ac17af3f 100644 --- a/src/crud/__init__.py +++ b/src/crud/__init__.py @@ -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 diff --git a/src/crud/deriver.py b/src/crud/deriver.py index 0852a479..770ba929 100644 --- a/src/crud/deriver.py +++ b/src/crud/deriver.py @@ -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, diff --git a/src/deriver/queue_manager.py b/src/deriver/queue_manager.py index b98c0ef6..493cddfa 100644 --- a/src/deriver/queue_manager.py +++ b/src/deriver/queue_manager.py @@ -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 Text, and_, delete, literal, 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, @@ -353,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, @@ -390,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(), @@ -406,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() diff --git a/src/dreamer/dream_due.py b/src/dreamer/dream_due.py new file mode 100644 index 00000000..08b68e2c --- /dev/null +++ b/src/dreamer/dream_due.py @@ -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 diff --git a/src/main.py b/src/main.py index a1ec9765..9a1d7e64 100644 --- a/src/main.py +++ b/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"]) diff --git a/src/reconciler/embed_now.py b/src/reconciler/embed_now.py index f76fa760..5308c062 100644 --- a/src/reconciler/embed_now.py +++ b/src/reconciler/embed_now.py @@ -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) diff --git a/src/reconciler/sync_vectors.py b/src/reconciler/sync_vectors.py index 1a8e99b5..b9b06418 100644 --- a/src/reconciler/sync_vectors.py +++ b/src/reconciler/sync_vectors.py @@ -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) diff --git a/src/routers/deriver_metrics.py b/src/routers/deriver_metrics.py new file mode 100644 index 00000000..547a5153 --- /dev/null +++ b/src/routers/deriver_metrics.py @@ -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, + } diff --git a/src/schemas/__init__.py b/src/schemas/__init__.py index 9f414583..0f93278a 100644 --- a/src/schemas/__init__.py +++ b/src/schemas/__init__.py @@ -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", diff --git a/src/schemas/internal.py b/src/schemas/internal.py index 2d299feb..f6399435 100644 --- a/src/schemas/internal.py +++ b/src/schemas/internal.py @@ -140,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.""" diff --git a/src/telemetry/prometheus/metrics.py b/src/telemetry/prometheus/metrics.py index cead893a..6859c7cd 100644 --- a/src/telemetry/prometheus/metrics.py +++ b/src/telemetry/prometheus/metrics.py @@ -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() diff --git a/tests/crud/test_deriver_metrics_query.py b/tests/crud/test_deriver_metrics_query.py new file mode 100644 index 00000000..a67c17a8 --- /dev/null +++ b/tests/crud/test_deriver_metrics_query.py @@ -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 diff --git a/tests/dreamer/test_dream_due.py b/tests/dreamer/test_dream_due.py new file mode 100644 index 00000000..dad75de6 --- /dev/null +++ b/tests/dreamer/test_dream_due.py @@ -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 diff --git a/tests/telemetry/test_metric_zero_init.py b/tests/telemetry/test_metric_zero_init.py index e69f50df..eb412287 100644 --- a/tests/telemetry/test_metric_zero_init.py +++ b/tests/telemetry/test_metric_zero_init.py @@ -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" # --------------------------------------------------------------------------- diff --git a/tests/test_deriver_metrics.py b/tests/test_deriver_metrics.py new file mode 100644 index 00000000..acd97330 --- /dev/null +++ b/tests/test_deriver_metrics.py @@ -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 From 7d5d6109f7ab2ba0368b6723e9eb213955b579e8 Mon Sep 17 00:00:00 2001 From: steven-ji Date: Thu, 3 Sep 2026 05:05:02 +0800 Subject: [PATCH 13/17] feat(docker): make API worker count configurable (#1088) * feat(docker): make API worker count configurable Add API_WORKERS with a single-worker default and document database pool sizing. Refs #1063 * fix(docker): address API worker review feedback --- .env.template | 3 +++ docker/entrypoint.sh | 2 +- docs/v3/contributing/self-hosting.mdx | 14 ++++++++++++++ 3 files changed, 18 insertions(+), 1 deletion(-) diff --git a/.env.template b/.env.template index 2f737dba..5c2b19c5 100644 --- a/.env.template +++ b/.env.template @@ -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 diff --git a/docker/entrypoint.sh b/docker/entrypoint.sh index bc8e3f37..a9f6ea78 100755 --- a/docker/entrypoint.sh +++ b/docker/entrypoint.sh @@ -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 diff --git a/docs/v3/contributing/self-hosting.mdx b/docs/v3/contributing/self-hosting.mdx index 02d361f1..4992530e 100644 --- a/docs/v3/contributing/self-hosting.mdx +++ b/docs/v3/contributing/self-hosting.mdx @@ -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` From a5fa8c39621b9b06ec1a90515e1103996aff521a Mon Sep 17 00:00:00 2001 From: Eugene Eisenstein Date: Wed, 2 Sep 2026 17:27:01 -0400 Subject: [PATCH 14/17] fix(dialectic): make workspace chat search before it answers (#1120) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The workspace agent's prefetch is an orientation overview — scale, active peers, their cards — not the corpus. `low` is the only reasoning level that explicitly sets TOOL_CHOICE="auto", so the model was free to skip tools entirely, and it did: every workspace_chat call in CI run 33662772219 made zero tool calls. It answered when the overview happened to carry the fact and otherwise wrote out the search it should have run, then asked the caller which option to take — at an endpoint with no caller to answer. Add a `_tool_choice` seam alongside `_select_tools` and override it on WorkspaceDialecticAgent to require a tool call. `execute_tool_loop` already relaxes "required"/"any" to "auto" after the first iteration, so this costs one search round rather than pinning the loop, and the model can still stop and synthesize. Any value a level configures other than None/"auto" passes through. The pair agent is unaffected: it prefetches the observations for its query and can legitimately answer from context alone. Also tell the workspace prompt it is non-interactive. It had "Do not narrate tool use" but never said the caller cannot reply, and three of the five traced responses ended in a menu of lookups. Unified subset goes 1/5 -> 5/5, and search_memory — the recall path that never once ran — now fires on 6 of 7 workspace queries. workspace_chat_scope is the notable one: its two not_contains assertions were passing vacuously because nothing was ever retrieved, and it now recalls the in-scope fact while still excluding the out-of-scope vault code. Co-authored-by: Claude Opus 5 (1M context) --- src/dialectic/core.py | 25 +++++++++++-- src/dialectic/prompts.py | 6 +++- src/dialectic/workspace.py | 26 +++++++++++++- tests/test_workspace_chat.py | 68 +++++++++++++++++++++++++++++++++--- 4 files changed, 116 insertions(+), 9 deletions(-) diff --git a/src/dialectic/core.py b/src/dialectic/core.py index 57964c87..95fe4cb5 100644 --- a/src/dialectic/core.py +++ b/src/dialectic/core.py @@ -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, diff --git a/src/dialectic/prompts.py b/src/dialectic/prompts.py index 5dfe6604..4d2fbc70 100644 --- a/src/dialectic/prompts.py +++ b/src/dialectic/prompts.py @@ -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 diff --git a/src/dialectic/workspace.py b/src/dialectic/workspace.py index 5383cd75..1a93d300 100644 --- a/src/dialectic/workspace.py +++ b/src/dialectic/workspace.py @@ -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, diff --git a/tests/test_workspace_chat.py b/tests/test_workspace_chat.py index daed20e8..75462b17 100644 --- a/tests/test_workspace_chat.py +++ b/tests/test_workspace_chat.py @@ -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] From 55a0519bd2e9db615bf4ce3d492558ad96b9fc47 Mon Sep 17 00:00:00 2001 From: steven-ji Date: Thu, 3 Sep 2026 05:31:46 +0800 Subject: [PATCH 15/17] feat(sdk): add per-call peer chat timeout (#1098) Forward optional timeout overrides through sync and async Peer.chat while retaining client-wide defaults. Refs #734 --- docs/v3/documentation/reference/sdk.mdx | 10 ++++++++ sdks/python/CHANGELOG.md | 6 +++++ sdks/python/src/honcho/aio.py | 10 +++++++- sdks/python/src/honcho/peer.py | 9 +++++++ tests/sdk/test_peer.py | 34 +++++++++++++++++++++++++ 5 files changed, 68 insertions(+), 1 deletion(-) diff --git a/docs/v3/documentation/reference/sdk.mdx b/docs/v3/documentation/reference/sdk.mdx index 432a4aab..66e08e89 100644 --- a/docs/v3/documentation/reference/sdk.mdx +++ b/docs/v3/documentation/reference/sdk.mdx @@ -276,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([ @@ -378,6 +381,13 @@ const bobConclusions = await alice.conclusionsOf("bob").list(); // Conclusions ``` +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: diff --git a/sdks/python/CHANGELOG.md b/sdks/python/CHANGELOG.md index 7fe5d8f7..2bea7442 100644 --- a/sdks/python/CHANGELOG.md +++ b/sdks/python/CHANGELOG.md @@ -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 diff --git a/sdks/python/src/honcho/aio.py b/sdks/python/src/honcho/aio.py index f5148ee6..3629551f 100644 --- a/sdks/python/src/honcho/aio.py +++ b/sdks/python/src/honcho/aio.py @@ -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: diff --git a/sdks/python/src/honcho/peer.py b/sdks/python/src/honcho/peer.py index 38edf269..bf6cf2d7 100644 --- a/sdks/python/src/honcho/peer.py +++ b/sdks/python/src/honcho/peer.py @@ -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: diff --git a/tests/sdk/test_peer.py b/tests/sdk/test_peer.py index 0c019d9b..d563d034 100644 --- a/tests/sdk/test_peer.py +++ b/tests/sdk/test_peer.py @@ -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], From b573a84806c1db317762bec65c69ce4ed808731a Mon Sep 17 00:00:00 2001 From: ajspig <46900795+ajspig@users.noreply.github.com> Date: Wed, 2 Sep 2026 17:56:09 -0400 Subject: [PATCH 16/17] Harness core (#1110) * chore: scaffold @honcho-ai/harness-core * feat(harness-core): resolve shared root config * feat(harness-core): send client identity headers on SDK requests * feat(harness-core): drop cloud vs custom api header * feat(harness-core): migrating v0 config to schema v1 on read * chore(harness-core): clean up * feat(config): describe oauth and host overrides in the v1 schema * chore: rename to harness-plugin-core * feat(harness-plugin-core): update telemetry headers on a live client. --- harness-plugin-core/.gitignore | 2 + harness-plugin-core/CHANGELOG.md | 10 + harness-plugin-core/README.md | 70 ++++++ harness-plugin-core/bun.lock | 25 ++ harness-plugin-core/package.json | 31 +++ harness-plugin-core/src/config.ts | 239 ++++++++++++++++++++ harness-plugin-core/src/index.ts | 29 +++ harness-plugin-core/src/telemetry.ts | 64 ++++++ harness-plugin-core/tests/config.test.ts | 71 ++++++ harness-plugin-core/tests/telemetry.test.ts | 56 +++++ harness-plugin-core/tsconfig.json | 13 ++ schemas/config/v1.json | 43 ++++ 12 files changed, 653 insertions(+) create mode 100644 harness-plugin-core/.gitignore create mode 100644 harness-plugin-core/CHANGELOG.md create mode 100644 harness-plugin-core/README.md create mode 100644 harness-plugin-core/bun.lock create mode 100644 harness-plugin-core/package.json create mode 100644 harness-plugin-core/src/config.ts create mode 100644 harness-plugin-core/src/index.ts create mode 100644 harness-plugin-core/src/telemetry.ts create mode 100644 harness-plugin-core/tests/config.test.ts create mode 100644 harness-plugin-core/tests/telemetry.test.ts create mode 100644 harness-plugin-core/tsconfig.json create mode 100644 schemas/config/v1.json diff --git a/harness-plugin-core/.gitignore b/harness-plugin-core/.gitignore new file mode 100644 index 00000000..f06235c4 --- /dev/null +++ b/harness-plugin-core/.gitignore @@ -0,0 +1,2 @@ +node_modules +dist diff --git a/harness-plugin-core/CHANGELOG.md b/harness-plugin-core/CHANGELOG.md new file mode 100644 index 00000000..ca098844 --- /dev/null +++ b/harness-plugin-core/CHANGELOG.md @@ -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] diff --git a/harness-plugin-core/README.md b/harness-plugin-core/README.md new file mode 100644 index 00000000..0b4525de --- /dev/null +++ b/harness-plugin-core/README.md @@ -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.` → 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' }) +``` diff --git a/harness-plugin-core/bun.lock b/harness-plugin-core/bun.lock new file mode 100644 index 00000000..1522ce60 --- /dev/null +++ b/harness-plugin-core/bun.lock @@ -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=="], + } +} diff --git a/harness-plugin-core/package.json b/harness-plugin-core/package.json new file mode 100644 index 00000000..13e59083 --- /dev/null +++ b/harness-plugin-core/package.json @@ -0,0 +1,31 @@ +{ + "name": "@honcho-ai/harness-plugin-core", + "version": "0.1.0", + "description": "Shared runtime for Honcho harness plugins", + "author": "Plastic Labs ", + "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" + } +} diff --git a/harness-plugin-core/src/config.ts b/harness-plugin-core/src/config.ts new file mode 100644 index 00000000..126cdb78 --- /dev/null +++ b/harness-plugin-core/src/config.ts @@ -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 +} + +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 { + return v !== null && typeof v === 'object' && !Array.isArray(v) +} + +/** Pre-schema files (no schemaVersion) → v1 keys. Host blocks included. */ +function migrate(file: unknown): Record { + if (!isObj(file)) return {} + const v = file.schemaVersion + if (typeof v === 'number' && v >= CONFIG_SCHEMA_VERSION) return { ...file } + const out: Record = { ...file } + const blocks: Record[] = [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 = 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(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 = { ...(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, 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(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 = {} + 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 | undefined, name: string): RootConfig { + if (!hosts || !isObj(hosts[name])) return {} + return pickRoot(hosts[name]) +} + +/** + * Highest wins: HONCHO_* env → overlay → hosts. → root → built-in. + */ +export function resolveConfig( + file: unknown, + opts: { host: string; env?: NodeJS.Dict; 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 = process.env): string { + return env.HONCHO_CONFIG_PATH || join(homedir(), '.honcho', 'config.json') +} + +export function loadConfig(opts: { + host: string + env?: NodeJS.Dict + 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 }) +} diff --git a/harness-plugin-core/src/index.ts b/harness-plugin-core/src/index.ts new file mode 100644 index 00000000..ab47a2b3 --- /dev/null +++ b/harness-plugin-core/src/index.ts @@ -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' diff --git a/harness-plugin-core/src/telemetry.ts b/harness-plugin-core/src/telemetry.ts new file mode 100644 index 00000000..eb8b2dc3 --- /dev/null +++ b/harness-plugin-core/src/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 +): Record { + const headers: Record = { [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, + id: TelemetryIdentity = {}, + extra?: Record +): Record { + return Object.assign(headers, telemetryHeaders(id, extra)) +} diff --git a/harness-plugin-core/tests/config.test.ts b/harness-plugin-core/tests/config.test.ts new file mode 100644 index 00000000..2c1b6a3e --- /dev/null +++ b/harness-plugin-core/tests/config.test.ts @@ -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') + }) +}) diff --git a/harness-plugin-core/tests/telemetry.test.ts b/harness-plugin-core/tests/telemetry.test.ts new file mode 100644 index 00000000..170c9e7a --- /dev/null +++ b/harness-plugin-core/tests/telemetry.test.ts @@ -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') + }) +}) diff --git a/harness-plugin-core/tsconfig.json b/harness-plugin-core/tsconfig.json new file mode 100644 index 00000000..96d10fea --- /dev/null +++ b/harness-plugin-core/tsconfig.json @@ -0,0 +1,13 @@ +{ + "compilerOptions": { + "target": "ES2022", + "module": "ESNext", + "moduleResolution": "bundler", + "allowImportingTsExtensions": true, + "noEmit": true, + "strict": true, + "skipLibCheck": true, + "types": ["node"] + }, + "include": ["src/**/*.ts"] +} diff --git a/schemas/config/v1.json b/schemas/config/v1.json new file mode 100644 index 00000000..be1384db --- /dev/null +++ b/schemas/config/v1.json @@ -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" } + } + } +} From 2ad56a4d715b015d744ea86f05a287069d264898 Mon Sep 17 00:00:00 2001 From: Aakash Kattelu Date: Wed, 2 Sep 2026 17:56:51 -0400 Subject: [PATCH 17/17] feat(mcp): add stdio host for local clients (#1102) * feat(mcp): add stdio host for local clients * feat(mcp): add Streamable HTTP host and image Long-lived HTTP entry for Docker and other process hosts, reusing createServer(). Dedicated mcp/Dockerfile; compose service beside api. * fix(mcp): stdio launcher cwd/silent and HTTP session bounds Pin bun --cwd so bunfig loads. Silence bun run. Require Bearer on HTTP. Idle-expire and cap in-memory MCP sessions. * fix(mcp): re-check bearer on established HTTP sessions Session lookup returned early without Authorization, so a missing or wrong token still 200'd after initialize. Bind each session to the init key and 401 on mismatch. * fix: nit cleaning claude command --------- Co-authored-by: ajspig --- docker-compose.yml.example | 28 ++++ mcp/.dockerignore | 7 + mcp/Dockerfile | 19 +++ mcp/README.md | 74 ++++++++-- mcp/bunfig.toml | 5 + mcp/package.json | 4 +- mcp/src/config.ts | 25 +++- mcp/src/http.test.ts | 56 ++++++++ mcp/src/http.ts | 281 +++++++++++++++++++++++++++++++++++++ mcp/src/stdio.ts | 30 ++++ mcp/tsconfig.json | 2 +- 11 files changed, 517 insertions(+), 14 deletions(-) create mode 100644 mcp/.dockerignore create mode 100644 mcp/Dockerfile create mode 100644 mcp/bunfig.toml create mode 100644 mcp/src/http.test.ts create mode 100644 mcp/src/http.ts create mode 100644 mcp/src/stdio.ts diff --git a/docker-compose.yml.example b/docker-compose.yml.example index ee201e6f..4bfae9f7 100644 --- a/docker-compose.yml.example +++ b/docker-compose.yml.example @@ -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 diff --git a/mcp/.dockerignore b/mcp/.dockerignore new file mode 100644 index 00000000..62b547a4 --- /dev/null +++ b/mcp/.dockerignore @@ -0,0 +1,7 @@ +node_modules +.wrangler +.dev.vars +.env +.env.* +*.log +dist diff --git a/mcp/Dockerfile b/mcp/Dockerfile new file mode 100644 index 00000000..60e0a0d1 --- /dev/null +++ b/mcp/Dockerfile @@ -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"] diff --git a/mcp/README.md b/mcp/README.md index 3bd24217..94abd88a 100644 --- a/mcp/README.md +++ b/mcp/README.md @@ -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 " +``` + +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 " diff --git a/mcp/bunfig.toml b/mcp/bunfig.toml new file mode 100644 index 00000000..9d1af97a --- /dev/null +++ b/mcp/bunfig.toml @@ -0,0 +1,5 @@ +[loader] +".md" = "text" + +[run] +silent = true diff --git a/mcp/package.json b/mcp/package.json index 900df5cf..231ec495 100644 --- a/mcp/package.json +++ b/mcp/package.json @@ -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" }, diff --git a/mcp/src/config.ts b/mcp/src/config.ts index ed70bbc1..ee8067e3 100644 --- a/mcp/src/config.ts +++ b/mcp/src/config.ts @@ -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, diff --git a/mcp/src/http.test.ts b/mcp/src/http.test.ts new file mode 100644 index 00000000..385f84eb --- /dev/null +++ b/mcp/src/http.test.ts @@ -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, 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); +}); diff --git a/mcp/src/http.ts b/mcp/src/http.ts new file mode 100644 index 00000000..a33bd4a8 --- /dev/null +++ b/mcp/src/http.ts @@ -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; +}; + +declare const Bun: { + serve(options: { + hostname: string; + port: number; + fetch(request: Request): Response | Promise; + }): { 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 = { + "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(); +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, +): 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 { + 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 { + 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}`); +} diff --git a/mcp/src/stdio.ts b/mcp/src/stdio.ts new file mode 100644 index 00000000..3a351651 --- /dev/null +++ b/mcp/src/stdio.ts @@ -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; + 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); +} diff --git a/mcp/tsconfig.json b/mcp/tsconfig.json index 8ca7bfae..8bebbfa9 100644 --- a/mcp/tsconfig.json +++ b/mcp/tsconfig.json @@ -10,5 +10,5 @@ "types": ["@cloudflare/workers-types"] }, "include": ["src/**/*.ts"], - "exclude": ["node_modules"] + "exclude": ["node_modules", "src/**/*.test.ts"] }