fix: only classify genuine oversize input as a token-limit error (#791)

Callers wrapped every ValueError from the embedding client in a
"exceeds maximum token limit" message, so provider and configuration
failures (dimension mismatch, empty response, upstream error) surfaced
to users as though their input were too long.

Add EmbeddingTokenLimitError, raised only by the pre-flight token checks
in embed() and simple_batch_embed(), and narrow the remaps in search.py,
agent_tools.py, document.py and representation.py to catch it. It
subclasses ValueError so existing broad handlers keep working.

Both simple_batch_embed() remap sites pass on_oversize="truncate" and so
could never raise a token-limit error at all; their handlers only ever
mislabelled provider failures.

Fixes #568

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

src/crud/document.py

@@ -18,7 +18,7 @@ from src.crud.collection import get_or_create_collection
 from src.crud.peer import get_peer, reject_scope_observed
 from src.crud.session import get_session
 from src.dependencies import tracked_db
-from src.embedding_client import embedding_client
+from src.embedding_client import EmbeddingTokenLimitError, embedding_client
 from src.exceptions import (
     ResourceNotFoundException,
     ValidationException,
@@ -362,7 +362,7 @@ async def query_documents(
     if embedding is None:
         try:
             embedding = await embedding_client.embed(query)
-        except ValueError as e:
+        except EmbeddingTokenLimitError as e:
             raise ValidationException(
                 "Query exceeds maximum token limit of "
                 + f"{settings.EMBEDDING.MAX_INPUT_TOKENS}."
@@ -987,7 +987,7 @@ async def create_observations(
         embeddings = await embedding_client.simple_batch_embed(
             contents, on_oversize="truncate"
         )
-    except ValueError as e:
+    except EmbeddingTokenLimitError as e:
         raise ValidationException(str(e)) from e
 
     # Create document objects and track embeddings for vector store

src/crud/representation.py

@@ -13,7 +13,7 @@ from src import crud, exceptions, models, schemas
 from src.config import settings
 from src.dependencies import tracked_db
 from src.dreamer.dream_scheduler import check_and_schedule_dream
-from src.embedding_client import embedding_client
+from src.embedding_client import EmbeddingTokenLimitError, embedding_client
 from src.schemas import ResolvedConfiguration
 from src.telemetry.events import EmbeddingCallPurpose
 from src.telemetry.logging import accumulate_metric
@@ -109,7 +109,7 @@ class RepresentationManager:
                 embeddings = await embedding_client.simple_batch_embed(
                     observation_texts, on_oversize="truncate"
                 )
-        except ValueError as e:
+        except EmbeddingTokenLimitError as e:
             raise exceptions.ValidationException(
                 "Observation content exceeds maximum token limit of "
                 + f"{settings.EMBEDDING.MAX_INPUT_TOKENS}."

src/embedding_client.py

@@ -159,6 +159,17 @@ def _publish_embedding_event(
         logger.debug("Failed to emit EmbeddingCallCompletedEvent", exc_info=True)
 
 
+class EmbeddingTokenLimitError(ValueError):
+    """Raised when input text genuinely exceeds the model's token limit.
+
+    Subclasses ``ValueError`` so existing broad handlers keep working, while
+    letting callers tell a real "content too long" condition apart from a
+    transient provider or configuration failure (dimension mismatch, empty
+    response, upstream error). Only the pre-flight token checks raise this;
+    provider failures keep raising plain ``ValueError``.
+    """
+
+
 class BatchItem(NamedTuple):
     """A single item in a batch with its metadata."""
 
@@ -272,7 +283,7 @@ class _EmbeddingClient:
         token_count = len(self.encoding.encode(query))
 
         if token_count > self.max_embedding_tokens:
-            raise ValueError(
+            raise EmbeddingTokenLimitError(
                 f"Query exceeds maximum token limit of {self.max_embedding_tokens} tokens (got {token_count} tokens)"
             )
 
@@ -358,8 +369,8 @@ class _EmbeddingClient:
             List of embedding vectors, one per input text (in order)
 
         Raises:
-            ValueError: If any text exceeds token limits and `on_oversize` is
-                ``"raise"``
+            EmbeddingTokenLimitError: If any text exceeds token limits and
+                `on_oversize` is ``"raise"``
         """
         if not texts:
             return []
@@ -380,7 +391,7 @@ class _EmbeddingClient:
                         tokens,
                     )
                 else:
-                    raise ValueError(
+                    raise EmbeddingTokenLimitError(
                         f"Text at index {idx} exceeds maximum token limit of "
                         + f"{self.max_embedding_tokens} tokens (got {len(token_ids)} tokens)"
                     )

src/utils/agent_tools.py

@@ -13,7 +13,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
 from src import crud, models, schemas
 from src.config import settings
 from src.dependencies import tracked_db
-from src.embedding_client import embedding_client
+from src.embedding_client import EmbeddingTokenLimitError, embedding_client
 from src.exceptions import ResourceNotFoundException
 from src.models import Document
 from src.schemas import ResolvedConfiguration
@@ -1884,11 +1884,15 @@ async def _handle_search_memory(
             parent_category=ctx.parent_category,
         ):
             query_embedding = await embedding_client.embed(query)
-    except ValueError:
+    except EmbeddingTokenLimitError:
         return (
             "ERROR: Query exceeds maximum token limit of "
             + f"{settings.EMBEDDING.MAX_INPUT_TOKENS}. Please use a shorter query."
         )
+    except ValueError as e:
+        # Provider/config failure, not an oversized query. Keep returning a
+        # string so the tool loop can continue, but don't blame the query.
+        return f"ERROR: Embedding the query failed: {e}"
 
     # Base telemetry metadata; results_count gets filled in below.
     search_meta: dict[str, Any] = {

src/utils/search.py

@@ -14,7 +14,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
 from src import models
 from src.config import settings
 from src.dependencies import tracked_db
-from src.embedding_client import embedding_client
+from src.embedding_client import EmbeddingTokenLimitError, embedding_client
 from src.exceptions import ValidationException
 from src.models import session_peers_table
 from src.telemetry.events import EmbeddingCallPurpose
@@ -388,7 +388,7 @@ async def search(
                 parent_category="api",
             ):
                 query_embedding = await embedding_client.embed(query)
-        except ValueError as e:
+        except EmbeddingTokenLimitError as e:
             raise ValidationException(
                 f"Query exceeds maximum token limit of {settings.EMBEDDING.MAX_INPUT_TOKENS}."
             ) from e

tests/llm/test_embedding_client.py

@@ -14,6 +14,7 @@ from src.config import (
 from src.embedding_client import (
     BatchItem,
     EmbeddingClient,
+    EmbeddingTokenLimitError,
     _EmbeddingClient,  # pyright: ignore[reportPrivateUsage]
 )
 
@@ -1173,3 +1174,70 @@ async def test_gemini_process_batch_wraps_contents_as_content_part(
     assert all(isinstance(c, genai_types.Content) for c in contents)
     assert contents[0].parts[0].text == "hello"
     assert contents[1].parts[0].text == "world"
+
+
+# --- Token-limit classification (issue #568) -------------------------------
+#
+# Only genuine "content too long" conditions may raise
+# EmbeddingTokenLimitError. Provider/config failures must stay plain
+# ValueError so callers don't rewrite them as token-limit errors.
+
+
+def test_embedding_token_limit_error_is_value_error() -> None:
+    """Subclassing ValueError keeps pre-existing broad handlers working."""
+    assert issubclass(EmbeddingTokenLimitError, ValueError)
+
+
+@pytest.mark.asyncio
+async def test_embed_raises_token_limit_error_before_calling_provider(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    client, fake_embeddings = _build_openai_client(
+        monkeypatch,
+        embedding=[0.1, 0.2],
+        model="text-embedding-3-small",
+        send_dimensions=False,
+        vector_dimensions=2,
+    )
+
+    with pytest.raises(EmbeddingTokenLimitError):
+        await client.embed("word " * 20_000)
+
+    assert fake_embeddings.calls == [], "provider must not be called on oversize input"
+
+
+@pytest.mark.asyncio
+async def test_simple_batch_embed_raises_token_limit_error_before_calling_provider(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    client, fake_embeddings = _build_openai_client(
+        monkeypatch,
+        embedding=[0.1, 0.2],
+        model="text-embedding-3-small",
+        send_dimensions=False,
+        vector_dimensions=2,
+    )
+
+    with pytest.raises(EmbeddingTokenLimitError):
+        await client.simple_batch_embed(["fine", "word " * 20_000])
+
+    assert fake_embeddings.calls == [], "provider must not be called on oversize input"
+
+
+@pytest.mark.asyncio
+async def test_provider_dimension_mismatch_is_not_a_token_limit_error(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A wrong-width vector is a provider/config fault, not an oversized input."""
+    client, _ = _build_openai_client(
+        monkeypatch,
+        embedding=[0.1, 0.2, 0.3],  # 3 wide, client expects 2
+        model="text-embedding-3-small",
+        send_dimensions=False,
+        vector_dimensions=2,
+    )
+
+    with pytest.raises(ValueError) as excinfo:
+        await client.embed("short query")
+
+    assert not isinstance(excinfo.value, EmbeddingTokenLimitError)