From f0f067d66b7e15c020caa334b3a6e5311e9483de Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=92=D0=BB=D0=B0=D0=B4=D0=B8=D0=BC=D0=B8=D1=80=20=D0=A8?= =?UTF-8?q?=D0=BC=D0=B0=D0=BD?= Date: Sat, 15 Aug 2026 18:34:22 +0300 Subject: [PATCH 1/6] A statusline that carries the reasoning, and a funnel that returns Five things, and three of them turned out to be one bug wearing different clothes: something took longer than a caller was willing to wait, and nothing on screen said so. **An agent statusline, and the reasoning behind it.** Always on screen above the composer, naming the call in flight rather than spinning; clicking it shows or hides the agent's reasoning. The strip it replaced appeared only while a turn ran and said "running ...", so the state worth reading at a glance was the one that came and went. Reasoning is a third kind of block in `TurnStream`, kept out of `text` -- the answer and the working are different claims -- and drawn whether or not it is switched on, with a class on the chat root deciding whether it is painted: a toggle that rebuilt the transcript would take its scroll position with it. Capturing it was not enough to have any: Opus 4.7+ defaults `display` to "omitted" and sends thinking blocks with a signature and no text, so `build_options` now asks for "summarized". **Interrupt, then a prompt, then nothing until you interrupt again.** One symptom, three causes, all ending with `busy` still set so the composer silently refused the next prompt. The SDK could refuse the interrupt (swallowed); the turn could ignore it (forever); or it could land late, on the turn issued after the one it was aimed at. So: the failure is reported, the turn is made to end -- `ui/tasks.py:cancel`'s escalation, asked first and inflicted second -- the client is rebuilt so no message from the stopped turn can end the next one, and the next turn waits for a pending interrupt. `drive_turn` now reports the SDK session id as it arrives, because the return value is not reached on the path that needs it most. **The funnel.** Four defects between it and a result, each hiding the next. `search_papers_by_relevance` takes `keyword`; Asta holds its event stream open and pings every 15s, which a per-read timeout cannot bound; its hits are wrapped under `result`, a key `_rows` did not know; and its `limit` must be <= 100, which `--no-expand` exceeded. Fixing all four leaves a corpus that answers a search in ~121s and takes ~283s to report its own backend refusing a connection -- twenty minutes of discovery that every caller kills first. So tier 1 now defaults to Papers with Code (`paperswithcode.co`, anonymous, 1-2s), with abstracts fetched from arXiv a hundred at a time because its search rows carry none. Asta stays behind `--tier1 asta`. A live run: 118 candidates in 3.1s, 99 abstracts in one request, 15 survivors. Stage 1 is now bounded twice over and says what it dropped: an endpoint that fails once is not asked again this run (it cost 283s to learn the same thing six times), and a wall clock stops discovery with what it has rather than being killed with nothing. It also prints progress to stderr, and the workspace runs its CLIs with `-u` -- everything that runs them reads a pipe, and Python block-buffers a pipe, so a ten-minute command delivered its output at exit and the tasks window, whose premise is that output is streamed as it arrives, showed an empty tail for the whole run. **The tasks window shows the agent's own calls.** Every capability here is reached by a Bash into `tools/`, so those are the other half of "what is running on this machine"; the transcript scrolls, and once it has, "is that still going" had no answer. Listed apart from the workspace's own tasks and carrying no STOP, because only interrupting the turn stops one. A call still running in a turn that already settled is reported as unfinished, not as running. **The session picker is the workspace's own menu.** It was the last Quasar control in an app whose stylesheet bypasses Quasar rather than overriding it -- and a `select` has one string per option, so "another window has this open" arrived as a fragment glued onto the title and looked exactly like a row that can be opened. `kit.Menu` is the shell's dialog helper, moved so the fourth of these could use it. Co-Authored-By: Claude Opus 5 --- README.md | 96 ++++++--- agent.py | 153 +++++++++++-- core/config.py | 61 +++++- core/http.py | 467 ++++++++++++++++++++++++++++++++++++---- tests/test_asta.py | 172 ++++++++++++++- tests/test_funnel.py | 195 ++++++++++++++++- tests/test_pwc.py | 301 ++++++++++++++++++++++++++ tests/test_streaming.py | 125 +++++++++++ tests/test_ui_shell.py | 58 ++++- tests/test_ui_tasks.py | 103 +++++++++ tests/test_ui_turns.py | 252 ++++++++++++++++++++++ tools/paper_search.py | 218 +++++++++++++++++-- ui/app.py | 159 ++++++++++++-- ui/kit.py | 87 ++++++++ ui/models.py | 96 ++++++++- ui/shell.py | 75 ++----- ui/state.py | 34 ++- ui/tasks.py | 8 + ui/tokens.py | 69 +++++- ui/windows/chat.py | 366 +++++++++++++++++++++++++------ ui/windows/tasks.py | 95 +++++++- 21 files changed, 2913 insertions(+), 277 deletions(-) create mode 100644 tests/test_pwc.py create mode 100644 tests/test_ui_turns.py diff --git a/README.md b/README.md index 7c9d2c0..07e1506 100644 --- a/README.md +++ b/README.md @@ -100,26 +100,47 @@ the only thing that forced a shell open beside the app on a fresh machine. The value goes down a pipe rather than in an argument — an argv is visible to anything that can list processes. -### Retrieval without an institutional email - -Tier 1 defaults to **Ai2's Asta** (`asta-tools.allen.ai`) rather than to the -Semantic Scholar REST API. It is the same corpus and it exposes the same -`snippet_search` the funnel is built around, but Semantic Scholar -[no longer issues API keys to free-domain email addresses][s2-keys], which -leaves a personal account on a shared anonymous pool that is rate limited often -enough that "no results" and "no key" are hard to tell apart. - -Asta is reached over streamable HTTP without adopting MCP as an architecture, -which is what §5 already said about that endpoint. Its key is optional and -raises limits rather than unlocking anything. Set `[retrieval] tier1` to `s2` -or `both` to change it, or pass `--tier1` per search. - -**The endpoint, the transport and the tool names are from Ai2's documentation; -the shape of each tool's result is not verified against the live service.** -`core/http.py:_rows` reads both plausible shapes and raises on anything else, -because a search that quietly returns nothing reads as "the literature has -nothing on this". - +### Retrieval without an institutional email, and without waiting + +Tier 1 defaults to **Papers with Code** (`paperswithcode.co/api/v1`) — the +catalogue as revived by Hugging Face, not the `.com` site Meta shut down. It is +anonymous and read-only: no account, no key, nothing to store. The endpoints are +those of [`huggingface/pwc-cli`][pwc-cli], which is the reference client for +this API. + +The choice is about latency, and the numbers are measured rather than assumed: + +| tier 1 | one search | key | text | +| --- | --- | --- | --- | +| `pwc` *(default)* | **1–2 s** | none | title only; abstracts fetched separately | +| `asta` | ~121 s, and ~283 s to report a backend failure | optional | full-text snippets | +| `s2` | fast when it answers | institutional addresses only | full-text snippets | + +Stage 0 turns one question into ~6 queries and each goes to two endpoints, so +Asta's per-call latency is twenty minutes of discovery before anything is +ranked — and every caller gives up first: the agent's own Bash tool backgrounds +a command at 120 s, and a shell `timeout` kills it. Semantic Scholar +[no longer issues API keys to free-domain email addresses][s2-keys], leaving a +personal account on a shared anonymous pool that is rate limited often enough +that "no results" and "no key" are hard to tell apart. + +**What the default costs.** Asta is the only one of the three with genuine +full-text snippets, which is what §5 designed stage-3 triage around. Under `pwc`, +triage reads the *abstract* instead — fetched from arXiv in one batched request +for the whole candidate pool (`core/http.py:arxiv_abstracts`), because nearly +every row is an arXiv paper and `id_list` takes a hundred ids at once. A +candidate that ends up with no abstract is reranked on its title, and the trace +says how many did. `pwc`'s expansion is also a *dense neighbour* rather than a +citation edge, and `neighbours` reports it as one rather than claiming the +citation graph §5 asks for. + +Set `[retrieval] tier1` to `asta`, `s2`, `both` (the two Semantic Scholar doors) +or `all`, or pass `--tier1` per search. Two wall clocks bound a run either way: +`request_deadline_s` caps one request, and `stage1_budget_s` caps the whole of +stage 1 — when it is spent the funnel keeps what it retrieved and records how +many queries it actually searched. + +[pwc-cli]: https://github.com/huggingface/pwc-cli [s2-keys]: https://www.semanticscholar.org/product/api ## Run @@ -258,12 +279,39 @@ autouse fixture in `tests/conftest.py` replaces `core.http._httpx`, because a suite that reaches the network does not fail, it *hangs*. Implemented but not exercised against a live service: the HF Jobs backend, the -SSH backend, Asta, Semantic Scholar, the OpenRouter reranker, Voyage embeddings, -the two Haiku funnel stages, Context7, ShinkaEvolve, and RepoWiki. They are -written against the documented interfaces and fail with actionable errors rather -than tracebacks, but a real credential and a real run are what will find the +SSH backend, Semantic Scholar, the OpenRouter reranker, Voyage embeddings, the +two Haiku funnel stages, Context7, ShinkaEvolve, and RepoWiki. They are written +against the documented interfaces and fail with actionable errors rather than +tracebacks, but a real credential and a real run are what will find the mismatches. +Tier-1 discovery is no longer in that list. A live run of the default funnel +returns 118 candidates from two rankings in 3.1 seconds, fills 99 abstracts in +one arXiv request, and hands 15 survivors to triage. + +Papers with Code and Asta are both exercised against the live services now, and +what that found is worth recording, because none of it could have been reasoned +out from the documentation: + +* **Asta's `search_papers_by_relevance` takes `keyword`** where `snippet_search` + takes `query`, so tier 1 lost that endpoint on every search — the server + answers in about a second with a validation error, which the slower endpoint's + latency hid. +* **Asta answers `tools/call` with an event stream it holds open**, pinging every + 15 seconds. `httpx`'s timeout is per socket read, so every ping reset it: a + buffered read waited for a close the server never promised, and discovery did + not fail, it never returned. The client now streams the response, stops at the + reply to its own request, and enforces a total deadline. +* **Asta wraps its hits under `result`** (singular), a key `_rows` did not know — + so even with the argument fixed, every hit was discarded as an unrecognised + envelope. And its `limit` must be ≤ 100, which `--no-expand` exceeded by + dividing the candidate ceiling across one query instead of six. +* **Papers with Code returns no abstract with a search row**, which is why + `arxiv_abstracts` exists. + +The *shapes* the remaining tools answer in are still read defensively, and an +unrecognised one still raises rather than returning an empty list. + **Two of [HANDOFF-2 §23](HANDOFF-2.md)'s open questions are now closed:** - **Context7's REST endpoints** (§23 item 2) are verified against the live API: diff --git a/agent.py b/agent.py index f0c6f8d..f49c3f6 100644 --- a/agent.py +++ b/agent.py @@ -22,9 +22,11 @@ import argparse import asyncio +import dataclasses import json import os import sys +import time from pathlib import Path from typing import Any @@ -73,22 +75,50 @@ def build_options(cfg: Any, *, permission_mode: str | None = None, resume: str | "PreToolUse": [sdk.HookMatcher(matcher="Bash", hooks=[hooks.pre_tool_use])], "Stop": [sdk.HookMatcher(hooks=[hooks.stop])], } - return sdk.ClaudeAgentOptions( - resume=resume, - model=cfg.model_for("research"), - system_prompt=system_prompt(), - allowed_tools=BUILTIN_TOOLS, - disallowed_tools=DENIED_TOOLS, - permission_mode=mode, - cwd=str(paths.root()), - hooks=hook_matchers, + options: dict[str, Any] = { + "resume": resume, + "model": cfg.model_for("research"), + "system_prompt": system_prompt(), + "allowed_tools": BUILTIN_TOOLS, + "disallowed_tools": DENIED_TOOLS, + "permission_mode": mode, + "cwd": str(paths.root()), + "hooks": hook_matchers, # Off by default in the SDK, and the default is why an answer used to # arrive in one lump: without it `receive_response` yields nothing until # a whole `AssistantMessage` is finished. With it the same turn also # emits `StreamEvent`s carrying token deltas. `TextStream` is what turns # the two into one transcript -- see the warning in its docstring. - include_partial_messages=True, - ) + "include_partial_messages": True, + } + options.update(thinking_option(cfg, sdk)) + return sdk.ClaudeAgentOptions(**options) + + +def thinking_option(cfg: Any, sdk: Any) -> dict[str, Any]: + """Ask for the reasoning as text, when the installed SDK can be asked. + + Capturing thinking blocks is not enough to *have* any: Opus 4.7+ defaults + `display` to "omitted" and sends them with a signature and no text. So the + chat window's reasoning switch had nothing to reveal no matter how correctly + the stream was read -- the bug was one flag away from the feature, and it + looked exactly like a toggle that did nothing. + + Feature-detected rather than assumed, for the same reason `agent.py probe` + exists: this option is newer than the permission mode and the SDK's shape has + changed between releases. An SDK without it gets the options it understands + and a session with no reasoning, which is what it would have had anyway. + """ + display = str(cfg.get("agent", "reasoning", "summarized")).lower() + if display not in ("summarized", "omitted"): + display = "summarized" + fields = {f.name for f in dataclasses.fields(sdk.ClaudeAgentOptions)} + if "thinking" not in fields: + return {} + # `adaptive` rather than a fixed budget: the model decides how much thinking + # a turn is worth, which is the right call for a session that ranges from + # "what is in the ledger" to a campaign design. + return {"thinking": {"type": "adaptive", "display": display}} def preflight_environment() -> dict[str, Any]: @@ -222,6 +252,7 @@ async def drive_turn( stream: Any, *, on_chunk: Any = None, + on_session_id: Any = None, session: str | None = None, ) -> dict[str, Any]: """One turn, for every surface that runs one. @@ -236,6 +267,13 @@ async def drive_turn( `on_chunk` is called with each newly-visible piece of text; the UI passes nothing because its renderer reads `stream.blocks` on a timer instead. + + `on_session_id` is called the moment the SDK names this conversation, and it + exists because the return value is not reached on every path this function + can take. A turn that is *interrupted* raises, so a caller reading the id + off the return value learned it only for turns that finished -- and an + interrupted turn is precisely the one after which the client is rebuilt, so + that was the case where losing the id cost the whole conversation. """ refusal = check_turn_budget() if refusal: @@ -258,6 +296,8 @@ async def drive_turn( # A resumed conversation can be given a new id, so the latest wins. candidate = getattr(message, "session_id", None) if isinstance(candidate, str) and candidate: + if candidate != sdk_session_id and on_session_id is not None: + on_session_id(candidate) sdk_session_id = candidate # The *last* usage seen, recorded once after the loop -- not one # record per message. `ResultMessage` arrives last and carries the @@ -324,6 +364,37 @@ def _delta_of(message: Any) -> str: return text if isinstance(text, str) else "" +def _thinking_delta_of(message: Any) -> str: + """The reasoning a partial-message stream event carries, if any. + + The mirror of `_delta_of`, and a separate function rather than a parameter on + it: the two feed different halves of the transcript, and the whole reason + `_delta_of` filters as hard as it does is that mixing them makes the stream + say something the settled message does not. + """ + event = getattr(message, "event", None) + if not isinstance(event, dict) or event.get("type") != "content_block_delta": + return "" + delta = event.get("delta") + if not isinstance(delta, dict) or delta.get("type") != "thinking_delta": + return "" + text = delta.get("thinking") + return text if isinstance(text, str) else "" + + +def _thinking_of(message: Any) -> str: + """The reasoning in a finished message. + + A `ThinkingBlock` carries `.thinking`, not `.text`, which is exactly why + `_text_of` misses it -- and why the reasoning needed a second pair of + accessors rather than a looser filter on the first. + """ + content = getattr(message, "content", None) + if not isinstance(content, list): + return "" + return "".join(getattr(b, "thinking", "") or "" for b in content) + + class TextStream: """One turn's visible text, assembled from deltas *and* finished messages. @@ -343,22 +414,29 @@ class TextStream: `feed` returns only the text that has not been shown yet, so a CLI can print its return value directly; `text` is the whole answer so far, for a UI that re-renders from it. + + The two accessors are parameters because the *reasoning* half of a turn + arrives the same way and has the same trap: `thinking_delta` events followed + by a `ThinkingBlock` containing all of them. One class, given the other pair + of accessors, is what keeps the no-duplication rule stated once. """ - def __init__(self) -> None: + def __init__(self, delta_of: Any = None, whole_of: Any = None) -> None: self.text = "" #: The tail of `text` contributed by deltas since the last finished #: message -- the part a finished message is entitled to overwrite. self._streamed = "" + self._delta_of = delta_of or _delta_of + self._whole_of = whole_of or _text_of def feed(self, message: Any) -> str: - delta = _delta_of(message) + delta = self._delta_of(message) if delta: self.text += delta self._streamed += delta return delta - text = _text_of(message) + text = self._whole_of(message) # A message with no text at all -- a tool result, a system message, the # final result -- must leave a half-streamed block alone. if not text: @@ -521,6 +599,13 @@ def tool_block(use: dict[str, Any]) -> dict[str, Any]: "rows": rows[:MAX_ROWS], "status": "running", "result": "", + # Wall clock rather than `time.monotonic`, because this is written into + # the session file and read back in another process, where a monotonic + # reading from a previous boot means nothing at all. What reads it is the + # tasks window, which reports how long the call in flight has been + # running -- the difference between a spinner and knowing a forty-minute + # job is still going. + "started": time.time(), } @@ -539,23 +624,42 @@ class TurnStream: happened. Text assembly is delegated to `TextStream` unchanged, including its rule that a finished message replaces the deltas that built it. + Reasoning is a third kind of block, kept beside the prose rather than folded + into it. It is *not* part of `text`: the answer and the working are different + claims, the transcript's `text` is what the agent said, and a UI that wants + the working can ask for the blocks. Which is exactly what the chat window's + statusline toggles. + `feed` returns what a CLI should print next; a UI re-renders from `blocks`. """ def __init__(self) -> None: self.blocks: list[dict[str, Any]] = [] self._text = TextStream() + self._think = TextStream(_thinking_delta_of, _thinking_of) #: The block the current run of prose is accumulating into, if open. self._open: dict[str, Any] | None = None + #: The same, for the current run of reasoning. + self._open_thought: dict[str, Any] | None = None #: Tool blocks by call id, so a result can find the call it answers. self._calls: dict[str, dict[str, Any]] = {} @property def text(self) -> str: - """The turn's prose, tool cards left out -- what a plain transcript says.""" + """The turn's prose, tool cards and reasoning left out -- what a plain + transcript says.""" return "".join(b["text"] for b in self.blocks if b["kind"] == "text") + @property + def thinking(self) -> str: + """The turn's reasoning, for a caller that wants only that half.""" + return "".join(b["text"] for b in self.blocks if b["kind"] == "thinking") + def feed(self, message: Any) -> str: + # Reasoning first, because that is the order the API sends it in: a + # message carrying both a `ThinkingBlock` and a `TextBlock` reasoned + # before it answered, and `blocks` is read top to bottom. + self._feed_thinking(message) printed = self._feed_text(message) for use in _tool_uses(message): block = tool_block(use) @@ -563,8 +667,12 @@ def feed(self, message: Any) -> str: self._calls[block["id"]] = block # A tool call ends the run of prose above it: whatever the agent says # next belongs *below* the card, because that is when it said it. + # The reasoning run ends with it, for the same reason -- interleaved + # thinking resumes *after* the call, not inside the block above it. self._text = TextStream() self._open = None + self._think = TextStream(_thinking_delta_of, _thinking_of) + self._open_thought = None printed += _tool_line(block) for result in _tool_results(message): block = self._calls.get(result["id"]) @@ -583,6 +691,8 @@ def note(self, text: str) -> None: if not text: return self._text = TextStream() + self._think = TextStream(_thinking_delta_of, _thinking_of) + self._open_thought = None self._open = {"kind": "text", "text": text} self.blocks.append(self._open) @@ -604,6 +714,19 @@ def _feed_text(self, message: Any) -> str: self._open["text"] = self._text.text return chunk + def _feed_thinking(self, message: Any) -> None: + """The same shape as `_feed_text`, over the reasoning accessors. + + Nothing is returned: `feed`'s return value is what a CLI prints, and the + reasoning is not that. It is a block a UI can choose to draw. + """ + self._think.feed(message) + if self._think.text: + if self._open_thought is None: + self._open_thought = {"kind": "thinking", "text": ""} + self.blocks.append(self._open_thought) + self._open_thought["text"] = self._think.text + def _tool_line(block: dict[str, Any]) -> str: subject = f" {block['title']}" if block["title"] else "" diff --git a/core/config.py b/core/config.py index 7a5f60f..36a3dbb 100644 --- a/core/config.py +++ b/core/config.py @@ -43,6 +43,13 @@ "retrieval": { "s2_base": "https://api.semanticscholar.org/graph/v1", "asta_base": "https://asta-tools.allen.ai/mcp/v1", + # Papers with Code, as revived by Hugging Face. The `.com` site Meta + # shut down is gone; this is the `.co` one, and its v1 API is anonymous, + # read-only and documented by `github.com/huggingface/pwc-cli`. + "pwc_base": "https://paperswithcode.co/api/v1", + # arXiv's Atom API, used for one thing: abstracts in bulk. See + # `arxiv_abstracts`. + "arxiv_base": "https://export.arxiv.org/api/query", "openrouter_base": "https://openrouter.ai/api/v1", "rerank_model": "voyageai/rerank-2.5", "embed_model": "voyage-4", @@ -64,17 +71,44 @@ "triage_top": 15, "cache_ttl_s": 604800, "request_timeout_s": 60, + # The ceiling on ONE request end to end, as distinct from + # `request_timeout_s`, which httpx applies per socket read. + # + # The distinction is the whole bug it exists to stop. Asta answers a + # `tools/call` with an event stream and holds it open while it works, + # sending `: ping` comments every 15s. Every ping resets the per-read + # timeout, so a 60s read timeout never fires no matter how long the + # server takes, and a buffered read of that stream waits for a close + # that may never come. A total deadline is the only thing that bounds + # it. 300s because a live `search_papers_by_relevance` measured 121s + # and a bound below the real latency is just an outage. + "request_deadline_s": 300, + # Wall clock for the whole of stage 1, as distinct from the per-request + # deadline above. Stage 0 turns one question into six queries and each + # goes to two endpoints, so a five-minute request deadline is a one-hour + # stage -- and the caller kills it long before the endpoints that work + # can contribute anything. When this is spent the funnel stops issuing + # tier-1 calls, keeps what it retrieved, and writes into the trace how + # many queries it actually searched. + "stage1_budget_s": 300, "min_request_interval_s": 1.1, # unauthenticated S2 is ~1 req/s - # Which tier-1 client does discovery: "asta", "s2", or "both". + # Which tier-1 client does discovery: "pwc", "asta", "s2", "both" + # (the two Semantic Scholar doors) or "all". # - # Asta by default because it is the one that *works*. Both reach the - # same Semantic Scholar corpus and both expose snippet search, but S2's - # own API stopped issuing keys to free-domain email addresses, leaving a - # personal account on the shared anonymous pool -- which is rate limited + # Papers with Code by default because it is the one that *answers*. + # Measured against the live services: pwc returns in 1-2s; Asta takes + # ~121s for a search and ~283s to report that its own backend refused a + # connection, and stage 0 multiplies that by six queries and two + # endpoints, so every caller gives up before discovery finishes. S2's + # own API stopped issuing keys to free-domain addresses, leaving a + # personal account on the shared anonymous pool, which is rate limited # often enough that "no results" and "no key" are hard to tell apart. - # Asta's key is optional. "both" is for comparing them, and it doubles - # the request count for candidates that mostly fuse back together. - "tier1": "asta", + # + # The trade is real and worth knowing: Asta is the only one of the three + # with genuine full-text snippets, which is what §5 designed stage-3 + # triage around. Under pwc, triage reads the abstract -- fetched from + # arXiv in one batched request, see `core/http.py:arxiv_abstracts`. + "tier1": "pwc", }, # HANDOFF-2 §18 listed the REST paths as unverified (§23 item 2). They are # now verified against the live API: `/api/v2/libs/search` returns @@ -139,6 +173,17 @@ "model": "claude-opus-5", "permission_mode": "dontAsk", "max_turns": 0, # 0 = unbounded + # Whether the reasoning arrives as *text*, which is what the chat + # window's statusline switches on and off. + # + # This is not the same question as whether the model thinks. Opus 4.7+ + # defaults `display` to "omitted" and sends thinking blocks with a + # signature and no text, so a client that captures reasoning correctly + # still has nothing to show -- which is exactly what a toggle over an + # empty transcript looks like. "summarized" is the only value that + # actually produces text; "omitted" is the SDK's own default and is here + # so turning the feature off is a config edit rather than a code one. + "reasoning": "summarized", }, "hosts": {}, } diff --git a/core/http.py b/core/http.py index 64d4434..ee6503e 100644 --- a/core/http.py +++ b/core/http.py @@ -51,7 +51,9 @@ def identifier_of(paper: dict[str, Any]) -> Any: return None -def candidate_id(identifier: Any, title: Any = "", text: Any = "") -> str | None: +def candidate_id( + identifier: Any, title: Any = "", text: Any = "", *, namespace: str = "s2" +) -> str | None: """The key the funnel fuses candidates on, or None when there is not one. Both tier-1 clients mint ids in one `s2:` namespace, deliberately: it is one @@ -70,13 +72,20 @@ def candidate_id(identifier: Any, title: Any = "", text: Any = "") -> str | None when there is not, which fuses genuine duplicates and separates genuine distinctions; and None when there is neither, because a hit with no id, no title and no text has nothing to rank and nothing to cite. + + `namespace` is what keeps that sharing honest once there is more than one + corpus. Asta and S2 share `s2:` because they are the same index and the same + ids; Papers with Code is a different catalogue with its own numbering, and + giving it the same prefix would fuse two unrelated papers whose ids happened + to collide. The digest fallback is namespaced too, so a title-only hit from + one corpus does not silently absorb a title-only hit from another. """ if identifier not in (None, ""): - return f"s2:{identifier}" + return f"{namespace}:{identifier}" material = f"{title or ''}\n{str(text or '')[:400]}".strip() if not material: return None - return "s2:t-" + hashlib.sha256(material.encode("utf-8")).hexdigest()[:16] + return f"{namespace}:t-" + hashlib.sha256(material.encode("utf-8")).hexdigest()[:16] def _httpx() -> Any: @@ -236,6 +245,287 @@ def neighbours(self, paper_id: str, *, direction: str = "citations", limit: int return out +# --------------------------------------------------------------------------- +# Papers with Code -- the corpus that answers in seconds +# --------------------------------------------------------------------------- +class PapersWithCode: + """Tier 1, and the default: the ML/AI catalogue behind `paperswithcode.co`. + + **Why this replaced Asta as the default.** Asta serves the right corpus and + is the only one of these with genuine full-text snippets, but measured + against the live service it answers a `search_papers_by_relevance` in ~121 + seconds and takes ~283 seconds to report that its own backend refused a + connection. Stage 0 turns one question into six queries and each goes to two + endpoints, so that is twenty minutes of discovery before anything is ranked, + and every caller -- the agent's Bash tool at 120s, a shell `timeout`, a + person -- gives up first. This answers in one to two seconds. A retriever + that returns is worth more than a better one that does not. + + **Two search modes, and they are genuinely different rankings.** `keyword` + is lexical and `semantic` is dense, which is exactly the pair `corpus.rrf` + exists to fuse -- so the two verbs the funnel already calls per query map + onto them without the funnel knowing anything changed. + + **What is given up, stated plainly.** Search rows carry no abstract, so §5's + "triage on ~500 words of the paper itself" becomes "triage on the abstract", + and the abstracts are fetched separately -- see `arxiv_abstracts`, which + fills the whole pool in one request because nearly every row here is an + arXiv paper. `related` is a *dense neighbour* rather than a citation edge, + and `neighbours` says so rather than presenting it as the citation graph + §5 asks for. + + Anonymous and read-only: no account, no key, nothing to store. The + endpoints and their parameters are read off `huggingface/pwc-cli`, which is + the reference client for this API. + """ + + #: The largest page this API will return, checked against the live service. + MAX_PAGE = 100 + + #: `funnel verb -> the API's search mode`. The funnel calls both per query + #: and fuses them, which is the whole point of there being two. + MODES = {"snippet_search": "semantic", "paper_search": "keyword"} + + def __init__(self, cfg: Config) -> None: + self.base = str(cfg.get("retrieval", "pwc_base")).rstrip("/") + self.timeout = float(cfg.get("retrieval", "request_timeout_s", 60)) + self.ttl = float(cfg.get("retrieval", "cache_ttl_s", 604800)) + self.interval = float(cfg.get("retrieval", "min_request_interval_s", 1.1)) + + def _get(self, path: str, params: dict[str, Any] | None = None) -> Any: + key = f"pwc:{self.base}:{path}:{json.dumps(params or {}, sort_keys=True)}" + hit = _cached(key, self.ttl) + if hit is not None: + return hit + _throttle("pwc", self.interval) + httpx = _httpx() + try: + resp = httpx.get( + f"{self.base}/{path.lstrip('/')}", + params=params, + headers={"Accept": "application/json", "User-Agent": "grad/1"}, + timeout=self.timeout, + ) + except Exception as exc: # noqa: BLE001 + raise UpstreamError( + f"Papers with Code request failed: {exc}", + fix="retry, or run with --local-only to search papers already ingested", + ) from exc + if resp.status_code == 429: + raise UpstreamError( + "Papers with Code rate-limited the request", + fix="wait and retry; this API is anonymous, so there is no key to raise it", + ) + if resp.status_code >= 400: + raise UpstreamError( + f"Papers with Code returned {resp.status_code}: {resp.text[:200]}", + fix=( + "the endpoint may have moved: check github.com/huggingface/pwc-cli and " + "set [retrieval] pwc_base in config/grad.toml" + ), + ) + try: + data = resp.json() + except Exception as exc: # noqa: BLE001 + raise UpstreamError( + f"Papers with Code returned a body that is not JSON: {resp.text[:200]}", + fix="retry; if it persists the API contract has changed", + ) from exc + _store(key, data) + return data + + def _search(self, query: str, mode: str, limit: int) -> list[dict[str, Any]]: + data = self._get( + "papers/search", + { + "q": query, + "page": 1, + "page_size": max(1, min(self.MAX_PAGE, int(limit))), + "mode": mode, + }, + ) + return _normalise_pwc(_pwc_rows(data, f"papers/search ({mode})"), f"pwc.{mode}") + + def snippet_search(self, query: str, limit: int = 20) -> list[dict[str, Any]]: + """The dense side. Named for the verb the funnel calls, not for what it + returns: there are no snippets here, and `arxiv_abstracts` is what makes + the candidates readable enough to rerank.""" + return self._search(query, self.MODES["snippet_search"], limit) + + def paper_search(self, query: str, limit: int = 20) -> list[dict[str, Any]]: + """The lexical side.""" + return self._search(query, self.MODES["paper_search"], limit) + + def neighbours( + self, paper_id: str, *, direction: str = "citations", limit: int = 20 + ) -> list[dict[str, Any]]: + """Related work, and **not** the citation graph. + + The rows this returns carry `provenance: "dense"` and a similarity + score: they are nearest neighbours in an embedding space, not papers + that cite or are cited by the seed. That still buys recall -- §5's point + is that expansion reaches papers no query string does -- but calling it + a citation edge would put a claim in the trace that the data does not + support, and a ledger entry's basis is the thing that must not be + overstated. The backward direction is refused for the same reason it is + on Asta: there is no endpoint for it, and answering it with the forward + one would double-count a single direction under two names. + """ + if direction != "citations": + return [] + data = self._get( + f"papers/{_quote(paper_id)}/related", + {"limit": max(1, min(self.MAX_PAGE, int(limit)))}, + ) + return _normalise_pwc(_pwc_rows(data, "related"), "pwc.related") + + +def _quote(value: Any) -> str: + """A path segment. Ids here are arXiv ids and integers, but they reach a URL + from a search result rather than from a constant.""" + from urllib.parse import quote # noqa: PLC0415 + + return quote(str(value), safe="._-") + + +def _pwc_rows(data: Any, what: str) -> list[dict[str, Any]]: + """The hits inside a response: `{"results": [...]}`, or a bare list. + + `related` answers with a list and `search` with an envelope, so both are + read. An unrecognised shape raises for the reason `_rows` does: `ok: true` + with no results reads as "the literature has nothing on this". + """ + if isinstance(data, list): + return [item for item in data if isinstance(item, dict)] + if isinstance(data, dict): + for key in ("results", "items", "data"): + if isinstance(data.get(key), list): + return [item for item in data[key] if isinstance(item, dict)] + raise UpstreamError( + f"Papers with Code's {what} returned no recognisable list of hits: " + f"{sorted(data)[:8] if isinstance(data, dict) else type(data).__name__}", + fix=( + "the API contract has changed -- compare it against " + "github.com/huggingface/pwc-cli and update core/http.py:_pwc_rows" + ), + ) + + +def _normalise_pwc(items: list[dict[str, Any]], source: str) -> list[dict[str, Any]]: + out = [] + for item in items: + row = _pwc_row(item, source=source) + if row is not None: + out.append(row) + return out + + +def _pwc_row(item: dict[str, Any], *, source: str) -> dict[str, Any] | None: + """One hit, in the vocabulary the funnel already fuses and reranks. + + The arXiv id leads, and that is deliberate: it is the field that survives + into `external`, that `paper_ingest` takes, and that `arxiv_abstracts` + fetches on -- so a candidate, its abstract and the ingest of its PDF all + name the same paper. The catalogue's own id is the fallback for the rows + that are not arXiv preprints. + """ + arxiv = str(item.get("arxiv_id") or "").strip() + identifier = arxiv or item.get("id") or item.get("route_identifier") + key = candidate_id(identifier, item.get("title"), "", namespace="pwc") + if key is None: + return None + published = str(item.get("published") or "") + return { + "id": key, + # What `neighbours` expands from: the API takes either, and the arXiv id + # is the one that means something outside this catalogue. + "paper_id": arxiv or str(item.get("id") or ""), + "title": item.get("title"), + "year": published[:4] or None, + # No full text here -- see the class docstring. `abstract` is filled in + # afterwards, and the reranker reads whichever of the two is present. + "snippet": "", + "abstract": str(item.get("abstract") or ""), + "section": "", + "source": source, + "citations": item.get("citation_count"), + "external": {"ArXiv": arxiv} if arxiv else {}, + "url": item.get("url_abs") or "", + } + + +# --------------------------------------------------------------------------- +# arXiv -- one request, every abstract +# --------------------------------------------------------------------------- +#: How many ids arXiv will take in one `id_list`. The whole reason to use this +#: rather than a per-paper lookup: a pool of a hundred candidates costs one +#: request instead of a hundred. +ARXIV_BATCH = 100 + + +def arxiv_abstracts( + arxiv_ids: Sequence[str], *, cfg: Config, timeout: float | None = None +) -> dict[str, str]: + """Abstracts for a batch of arXiv ids, as `{id: abstract}`. + + This exists because the fast corpus does not carry abstracts in its search + results and the reranker and stage-3 triage both read them: a candidate pool + of titles alone is a measurably worse funnel. Fetching them one at a time + would cost a request per candidate and undo the reason the corpus was + changed, so this uses `id_list`, which takes a hundred at once. + + Never raises. A missing abstract is a candidate that ranks on its title, + which is what would have happened anyway -- so a failure here degrades the + ranking rather than the run, and the caller records it as a warning. + """ + ids = [str(i).strip() for i in arxiv_ids if str(i or "").strip()][:ARXIV_BATCH] + if not ids: + return {} + key = f"arxiv:abstracts:{json.dumps(sorted(ids))}" + hit = _cached(key, float(cfg.get("retrieval", "cache_ttl_s", 604800))) + if isinstance(hit, dict): + return hit + _throttle("arxiv", max(3.0, float(cfg.get("retrieval", "min_request_interval_s", 1.1)))) + httpx = _httpx() + try: + resp = httpx.get( + str(cfg.get("retrieval", "arxiv_base")), + params={"id_list": ",".join(ids), "max_results": len(ids)}, + headers={"User-Agent": "grad/1"}, + timeout=timeout or float(cfg.get("retrieval", "request_timeout_s", 60)), + ) + resp.raise_for_status() + out = _parse_arxiv_atom(resp.text) + except Exception: # noqa: BLE001 - see the docstring: this degrades, it does not fail + return {} + _store(key, out) + return out + + +def _parse_arxiv_atom(xml: str) -> dict[str, str]: + """`{bare arxiv id: abstract}` out of an Atom feed. + + The id in the feed is a URL with a version suffix (`.../abs/1706.03762v7`) + and the id asked for has neither, so it is reduced to the bare form -- the + caller looks its candidates up by what Papers with Code gave it. + """ + import xml.etree.ElementTree as ET # noqa: PLC0415 + + namespace = {"atom": "http://www.w3.org/2005/Atom"} + try: + root = ET.fromstring(xml) + except ET.ParseError: + return {} + out: dict[str, str] = {} + for entry in root.findall("atom:entry", namespace): + raw = (entry.findtext("atom:id", "", namespace) or "").rsplit("/", 1)[-1] + bare = raw.split("v")[0] if raw else "" + summary = " ".join((entry.findtext("atom:summary", "", namespace) or "").split()) + if bare and summary: + out[bare] = summary + return out + + # --------------------------------------------------------------------------- # Asta -- the same corpus, through a door that opens # --------------------------------------------------------------------------- @@ -244,6 +534,18 @@ def neighbours(self, paper_id: str, *, direction: str = "citations", limit: int #: `initialize` result and this follows it. MCP_PROTOCOL_VERSION = "2025-06-18" +#: The largest `limit` Asta's tools accept, from the service: *"The limit +#: parameter must be between 1 and 100 inclusive."* +#: +#: Clamped here rather than left to callers because it is the service's +#: constraint, not the funnel's. `paper_search.py` divides its candidate ceiling +#: across the expanded queries, so skipping stage 0 -- one query instead of six +#: -- asks for six times as many per call, and the funnel's own `--no-expand` +#: path therefore refused *every* tier-1 call with a validation error. Asking +#: for the most the service will give is the honest reading of "as many as you +#: can"; raising `--candidates` should not be a way to get zero results. +MAX_LIMIT = 100 + class Asta: """Ai2's scientific corpus over MCP, at `asta-tools.allen.ai`. @@ -278,6 +580,10 @@ class Asta: def __init__(self, cfg: Config) -> None: self.base = str(cfg.get("retrieval", "asta_base")).rstrip("/") self.timeout = float(cfg.get("retrieval", "request_timeout_s", 60)) + #: Total wall clock for one request, enforced here rather than by httpx. + #: See `_post`: the per-read timeout above cannot bound a stream that is + #: being kept alive. + self.deadline = float(cfg.get("retrieval", "request_deadline_s", 300)) self.ttl = float(cfg.get("retrieval", "cache_ttl_s", 604800)) self.interval = float(cfg.get("retrieval", "min_request_interval_s", 1.1)) try: @@ -313,23 +619,57 @@ def _headers(self) -> dict[str, str]: return headers def _post(self, body: dict[str, Any]) -> Any: + """One JSON-RPC message out, one back -- under a deadline this enforces. + + **The response is streamed rather than buffered, and that is the whole + fix.** Asta answers `tools/call` with an event stream and holds it open + while it works, sending `: ping` comments every 15 seconds. `httpx`'s + `timeout` is per socket read, so every ping reset it: a buffered + `httpx.post` waited for a close that the server had no obligation to + send, the read timeout could never fire no matter how long it took, and + the funnel's first tier-1 call simply never returned. Discovery was + unreachable, and the failure had no error to go with it. + + Two things bound it now. `_mcp_payload` stops at the reply to *this* + request rather than at the end of the stream -- the answer is what was + asked for, and reading past it is waiting for a close nobody promised -- + and `deadline` caps the whole exchange in case the reply never comes. + """ _throttle("asta", self.interval) httpx = _httpx() + deadline = time.monotonic() + self.deadline try: - resp = httpx.post(self.base, json=body, headers=self._headers(), timeout=self.timeout) + with httpx.stream( + "POST", self.base, json=body, headers=self._headers(), timeout=self.timeout + ) as resp: + # Assigned on `initialize` and echoed from then on. A server that + # does not use sessions simply never sends it. + session = resp.headers.get("mcp-session-id") + if session: + self._session = session + if resp.status_code >= 400: + self._refuse(resp) + # A notification gets 202 Accepted and an empty body; there is + # nothing to parse and nothing to wait for. + if resp.status_code == 202: + return None + return _mcp_payload(resp, request_id=body.get("id"), deadline=deadline) + except UpstreamError: + # Already the shaped refusal, with the fix that belongs to it. Left + # alone rather than re-wrapped as a transport failure -- "Asta rate + # limited the request" and "Asta request failed" send you to two + # different places. + raise except Exception as exc: # noqa: BLE001 raise UpstreamError( f"Asta request failed: {exc}", fix="retry, or run with --local-only to search papers already ingested", ) from exc - # Assigned on `initialize` and echoed from then on. A server that does - # not use sessions simply never sends it. - session = resp.headers.get("mcp-session-id") - if session: - self._session = session - - if resp.status_code == 401 or resp.status_code == 403: + def _refuse(self, resp: Any) -> None: + """Turn a 4xx/5xx into the refusal that says what to do about it.""" + resp.read() + if resp.status_code in (401, 403): raise UpstreamError( f"Asta rejected the request ({resp.status_code})", fix=( @@ -346,19 +686,13 @@ def _post(self, body: dict[str, Any]) -> Any: f"python -m tools.jobs credential set {credentials.ASTA_KEY}" ), ) - if resp.status_code >= 400: - raise UpstreamError( - f"Asta returned {resp.status_code}: {resp.text[:200]}", - fix=( - "the endpoint may have moved: check allenai.org/asta/resources/mcp and " - "set [retrieval] asta_base in config/grad.toml" - ), - ) - # A notification gets 202 Accepted and an empty body; there is nothing - # to parse and nothing to wait for. - if resp.status_code == 202 or not (resp.content or b"").strip(): - return None - return _mcp_payload(resp) + raise UpstreamError( + f"Asta returned {resp.status_code}: {resp.text[:200]}", + fix=( + "the endpoint may have moved: check allenai.org/asta/resources/mcp and " + "set [retrieval] asta_base in config/grad.toml" + ), + ) def _call(self, method: str, params: dict[str, Any] | None = None) -> Any: self._id += 1 @@ -433,11 +767,19 @@ def tool(self, name: str, arguments: dict[str, Any]) -> Any: def snippet_search(self, query: str, limit: int = 20) -> list[dict[str, Any]]: """Full-text excerpts. The reason to prefer this over a metadata search: ~500 words of the paper itself is what stage 3 triages on.""" - data = self.tool("snippet_search", {"query": query, "limit": limit}) + data = self.tool("snippet_search", {"query": query, "limit": _limit(limit)}) return _normalise(_rows(data, "snippet_search"), "asta.snippet") def paper_search(self, query: str, limit: int = 20) -> list[dict[str, Any]]: - data = self.tool("search_papers_by_relevance", {"query": query, "limit": limit}) + # `keyword`, not `query`. The two tools disagree: `snippet_search` takes + # `query` and this one takes `keyword`, and sending the wrong one is not + # a soft failure -- the server answers in ~1s with a pydantic validation + # error ("keyword Field required"), which `tool` raises as an + # UpstreamError, so tier 1 lost this endpoint entirely on every search + # while snippet_search's slowness hid it. Checked against tools/list. + data = self.tool( + "search_papers_by_relevance", {"keyword": query, "limit": _limit(limit)} + ) return _normalise(_rows(data, "search_papers_by_relevance"), "asta.paper") def neighbours( @@ -453,20 +795,41 @@ def neighbours( """ if direction != "citations": return [] - data = self.tool("get_citations", {"paper_id": paper_id, "limit": limit}) + data = self.tool("get_citations", {"paper_id": paper_id, "limit": _limit(limit)}) return _normalise(_rows(data, "get_citations"), "asta.citations") -def _mcp_payload(resp: Any) -> Any: +def _limit(value: Any) -> int: + """A `limit` the service will accept. See `MAX_LIMIT`.""" + try: + wanted = int(value) + except (TypeError, ValueError): + return 20 + return max(1, min(MAX_LIMIT, wanted)) + + +def _mcp_payload(resp: Any, *, request_id: Any = None, deadline: float | None = None) -> Any: """One JSON-RPC message out of a streamable-HTTP response. The same request may be answered with `application/json` or with an SSE - stream, at the server's discretion, so both are handled. For a stream the - *last* `data:` frame carrying a result is taken: progress notifications - share the channel with the answer. + stream, at the server's discretion, so both are handled. + + For a stream this **stops at the reply to `request_id`** rather than reading + to the end. That is not an optimisation: the server may keep the stream open + after answering, pinging every 15 seconds, and a reader that waits for the + close waits forever -- which is exactly how a 60-second read timeout failed + to bound a call that never returned. The reply is the answer; there is + nothing after it worth waiting for. + + Progress notifications share the channel, so frames are filtered to those + carrying a `result` or an `error`, and the last one seen is the fallback for + a server that answers without echoing the id. """ content_type = (resp.headers.get("content-type") or "").lower() if "text/event-stream" not in content_type: + resp.read() + if not (resp.content or b"").strip(): + return None try: return resp.json() except Exception as exc: # noqa: BLE001 @@ -476,15 +839,27 @@ def _mcp_payload(resp: Any) -> Any: ) from exc answer: Any = None - for line in resp.text.splitlines(): + for line in resp.iter_lines(): + if deadline is not None and time.monotonic() > deadline: + raise UpstreamError( + "Asta held the stream open past its deadline without answering", + fix=( + "retry; raise [retrieval] request_deadline_s in config/grad.toml if the " + "corpus is genuinely this slow, or use --local-only meanwhile" + ), + ) + line = str(line).rstrip("\r") if not line.startswith("data:"): continue try: frame = json.loads(line[5:].strip()) except json.JSONDecodeError: continue - if isinstance(frame, dict) and ("result" in frame or "error" in frame): - answer = frame + if not isinstance(frame, dict) or not ("result" in frame or "error" in frame): + continue + answer = frame + if request_id is None or frame.get("id") == request_id: + return answer if answer is None: raise UpstreamError( "Asta's event stream carried no result", @@ -521,18 +896,32 @@ def _mcp_result(result: Any) -> Any: return text +#: The keys a tool's payload may wrap its hits under, in the order they are +#: tried. `result` is first because it is the one a live call actually returns +#: -- `search_papers_by_relevance` answers `{"result": [{"paperId": …, "title": +#: …}]}` -- and its absence is why tier 1 found nothing even once the argument +#: name was right: every hit was thrown away as an unrecognised envelope, and +#: the funnel raised rather than returning them. +#: +#: Singular, and *not* the JSON-RPC `result` field. By the time this runs, `tool` +#: has already unwrapped the RPC envelope and the tool's own content block, so +#: this is the payload's own key -- they are only spelled the same. +ROW_KEYS = ("result", "data", "results", "snippets", "papers", "citations", "items") + + def _rows(data: Any, tool_name: str) -> list[dict[str, Any]]: """The list of hits inside a tool's payload, whatever it is wrapped in. - Unverified against the live service, so this reads the S2 REST shape and the - obvious flattenings of it. An unrecognised payload raises: a search that - quietly returns nothing reads as "the literature has nothing on this", and - that is a conclusion nobody should draw from a schema change. + Partly verified against the live service now: `search_papers_by_relevance` + uses `result`. The rest are the S2 REST shape and the obvious flattenings of + it, still unverified. An unrecognised payload raises: a search that quietly + returns nothing reads as "the literature has nothing on this", and that is a + conclusion nobody should draw from a schema change. """ if isinstance(data, list): candidates = data elif isinstance(data, dict): - for key in ("data", "results", "snippets", "papers", "citations", "items"): + for key in ROW_KEYS: if isinstance(data.get(key), list): candidates = data[key] break diff --git a/tests/test_asta.py b/tests/test_asta.py index 8c9ebe8..f856040 100644 --- a/tests/test_asta.py +++ b/tests/test_asta.py @@ -28,17 +28,48 @@ class FakeResponse: + """A streamed response, because that is what the client now opens. + + `read()` and `iter_lines()` are the two halves of `httpx`'s streaming API the + client uses, and they are not interchangeable: the JSON path buffers, and the + event-stream path must *not*, because the server keeps the stream open after + answering. `lines` counts how far a test's reader actually got, which is what + makes "it stopped at the reply" assertable. + """ + def __init__(self, status_code=200, payload=None, text="", content_type="application/json", - headers=None): + headers=None, hang=False): self.status_code = status_code self._payload = payload if payload is not None else {} self.text = text or (json.dumps(self._payload) if payload is not None else "") self.headers = {"content-type": content_type, **(headers or {})} + self.content = b"" + #: A server that never closes the stream: `iter_lines` keeps yielding + #: keepalive comments after the body, the way Asta's does. + self.hang = hang + self.lines = 0 + + def read(self): self.content = self.text.encode() + return self.content + + def iter_lines(self): + for line in self.text.splitlines(): + self.lines += 1 + yield line + while self.hang: + self.lines += 1 + yield ": ping" def json(self): return self._payload + def __enter__(self): + return self + + def __exit__(self, *_): + return False + def rpc(result) -> dict: return {"jsonrpc": "2.0", "id": 1, "result": result} @@ -58,8 +89,8 @@ def transport(monkeypatch): class FakeHttpx: @staticmethod - def post(url, json=None, headers=None, timeout=None): - posts.append({"url": url, "body": json, "headers": headers}) + def stream(method, url, json=None, headers=None, timeout=None): + posts.append({"url": url, "method": method, "body": json, "headers": headers}) if queue: return queue.pop(0) return FakeResponse(payload=rpc({"protocolVersion": "2025-06-18"})) @@ -438,12 +469,14 @@ def test_the_tier_one_selector_builds_the_clients_the_config_names(workspace): from tools import paper_search cfg = config_mod.load(reload=True) + assert [n for n, _ in paper_search.tier1_clients(cfg, "pwc")] == ["pwc"] assert [n for n, _ in paper_search.tier1_clients(cfg, "asta")] == ["asta"] assert [n for n, _ in paper_search.tier1_clients(cfg, "s2")] == ["s2"] assert [n for n, _ in paper_search.tier1_clients(cfg, "both")] == ["asta", "s2"] + assert [n for n, _ in paper_search.tier1_clients(cfg, "all")] == ["pwc", "asta", "s2"] assert paper_search.tier1_clients(cfg, "none") == [] # The default, when nothing overrides it. - assert [n for n, _ in paper_search.tier1_clients(cfg)] == ["asta"] + assert [n for n, _ in paper_search.tier1_clients(cfg)] == ["pwc"] def test_an_unknown_tier_one_source_lists_the_real_ones(workspace): @@ -452,4 +485,133 @@ def test_an_unknown_tier_one_source_lists_the_real_ones(workspace): with pytest.raises(UsageError) as exc: paper_search.tier1_clients(config_mod.load(reload=True), "scholar") - assert "asta" in (exc.value.fix or "") + assert "pwc" in (exc.value.fix or "") + + +# --------------------------------------------------------------------------- +# the stream that is held open -- why the funnel's first call never returned +# --------------------------------------------------------------------------- +def sse(*frames: dict) -> str: + return "".join(f"event: message\ndata: {json.dumps(f)}\n\n" for f in frames) + + +def handshake(queue) -> None: + queue.append(FakeResponse(payload=rpc({"protocolVersion": "2025-06-18"}))) + queue.append(FakeResponse(status_code=202, text="")) + + +def test_the_read_stops_at_the_reply_rather_than_at_the_close(workspace, transport): + """The bug this closes. Asta answers `tools/call` with an event stream and + keeps it open afterwards, pinging every 15s. `httpx`'s timeout is per read, + so every ping reset it and a buffered read waited for a close the server + never promised -- the call did not fail, it never returned, and tier 1 was + simply unreachable with no error to go with it.""" + _, queue = transport + handshake(queue) + # id 2: initialize was 1, and the notification between them carries none. + answer = FakeResponse( + text=sse( + {"jsonrpc": "2.0", "method": "notifications/progress", "params": {}}, + rpc(tool_result({"data": [{"paperId": "p1", "title": "Attention"}]})) | {"id": 2}, + ), + content_type="text/event-stream", + hang=True, + ) + queue.append(answer) + + rows = client().snippet_search("attention") + assert [r["title"] for r in rows] == ["Attention"] + assert answer.lines <= 8, "it kept reading past the answer it already had" + + +def test_a_stream_that_never_answers_is_bounded_by_the_deadline(workspace, transport): + """`request_timeout_s` is per socket read and a keepalive resets it, so the + only thing that can bound this is a total deadline. A search that hangs is + worse than one that fails: nothing points at its own cause.""" + from core import paths + + config = paths.root() / "config" / "grad.toml" + config.parent.mkdir(parents=True, exist_ok=True) + config.write_text("[retrieval]\nrequest_deadline_s = -1\n", encoding="utf-8") + + _, queue = transport + handshake(queue) + queue.append(FakeResponse(text="", content_type="text/event-stream", hang=True)) + + with pytest.raises(UpstreamError, match="deadline"): + client().snippet_search("attention") + + +def test_the_deadline_message_names_the_key_that_moves_it(workspace, transport): + from core import paths + + config = paths.root() / "config" / "grad.toml" + config.parent.mkdir(parents=True, exist_ok=True) + config.write_text("[retrieval]\nrequest_deadline_s = -1\n", encoding="utf-8") + + _, queue = transport + handshake(queue) + queue.append(FakeResponse(text="", content_type="text/event-stream", hang=True)) + with pytest.raises(UpstreamError) as exc: + client().snippet_search("attention") + assert "request_deadline_s" in (exc.value.fix or "") + + +def test_relevance_search_sends_the_argument_that_tool_actually_takes(workspace, transport): + """`snippet_search` takes `query` and this one takes `keyword`, and sending + the wrong one is not a soft failure: the server answers in ~1s with a + pydantic validation error, so tier 1 lost this endpoint on every search + while snippet_search's slowness hid it.""" + posts, queue = transport + handshake(queue) + queue.append(FakeResponse(payload=rpc(tool_result({"data": []})))) + + client().paper_search("efficient optimizers", limit=7) + arguments = posts[-1]["body"]["params"]["arguments"] + assert arguments == {"keyword": "efficient optimizers", "limit": 7} + + +def test_the_envelope_a_live_call_actually_uses_is_read(workspace, transport): + """`search_papers_by_relevance` answers `{"result": [...]}` -- singular, and + a key this did not know. Every hit was thrown away as an unrecognised + envelope, so tier 1 found nothing even once the argument name was right, and + the funnel raised rather than returning them. + + Not the JSON-RPC `result`: by the time `_rows` runs, the RPC envelope and the + tool's content block are both already unwrapped. They are only spelled the + same.""" + rows = call_with(workspace, transport, {"result": [ + {"paperId": "sha-1", "title": "Memory Efficient Optimizers with 4-bit States"}, + ]}) + assert [r["id"] for r in rows] == ["s2:sha-1"] + assert rows[0]["title"].startswith("Memory Efficient") + + +def test_a_still_unrecognised_envelope_names_the_keys_it_knows(workspace, transport): + from core import http as http_mod + + assert "result" in http_mod.ROW_KEYS + with pytest.raises(UpstreamError, match="no recognisable list"): + call_with(workspace, transport, {"payload": {"hits": []}}) + + +def test_a_limit_above_what_the_service_accepts_is_clamped(workspace, transport): + """*"The limit parameter must be between 1 and 100 inclusive."* The funnel + divides its candidate ceiling across the expanded queries, so skipping stage + 0 asks for six times as many per call -- and `--no-expand` therefore refused + every tier-1 call with a validation error. Raising `--candidates` should not + be a way to get zero results.""" + posts, queue = transport + handshake(queue) + queue.append(FakeResponse(payload=rpc(tool_result({"result": []})))) + + client().paper_search("efficient optimizers", limit=150) + assert posts[-1]["body"]["params"]["arguments"]["limit"] == http.MAX_LIMIT + + +def test_a_nonsense_limit_becomes_a_usable_one_rather_than_a_refusal(workspace, transport): + posts, queue = transport + handshake(queue) + queue.append(FakeResponse(payload=rpc(tool_result({"result": []})))) + client().snippet_search("attention", limit=0) + assert posts[-1]["body"]["params"]["arguments"]["limit"] == 1 diff --git a/tests/test_funnel.py b/tests/test_funnel.py index 52858cf..d1b8306 100644 --- a/tests/test_funnel.py +++ b/tests/test_funnel.py @@ -135,7 +135,8 @@ def test_a_healthy_stage_still_returns_its_payload(monkeypatch): "source, client, expected_fix", [ # Each door's dead end points at the one that is not a dead end. - ("s2", "SemanticScholar", "--tier1 asta"), + ("pwc", "PapersWithCode", "--tier1 asta"), + ("s2", "SemanticScholar", "--tier1 pwc"), ("asta", "Asta", "credential set asta_api_key"), ], ) @@ -191,7 +192,7 @@ def test_an_empty_run_still_writes_the_trace_the_funnel_view_reads( from tools import paper_search # The default source, so this covers the path a real run takes. - monkeypatch.setattr(http, "Asta", lambda cfg: _RateLimited()) + monkeypatch.setattr(http, "PapersWithCode", lambda cfg: _RateLimited()) paper_search.cli.run( ["search", "efficient optimizers", "--no-expand", "--no-triage", "--no-local", "--json"] ) @@ -202,7 +203,7 @@ def test_an_empty_run_still_writes_the_trace_the_funnel_view_reads( written = json.loads(traces[0].read_text(encoding="utf-8")) assert written["stages"]["1_retrieve"] == {"rankings": 0, "candidates": 0} # Which client was asked, so "no results" can be read against who was down. - assert written["stages"]["1_sources"] == ["asta"] + assert written["stages"]["1_sources"] == ["pwc"] assert any("rate-limited" in w for w in written["warnings"]) @@ -311,3 +312,191 @@ def test_the_fix_the_error_prints_is_a_command_that_exists(): assert credentials.CLAUDE_TOKEN in CREDENTIAL_NAMES assert credentials.CLAUDE_TOKEN in credentials.status() assert credentials.CLAUDE_TOKEN in haiku.AUTH_FIX + + +# --------------------------------------------------------------------------- +# stage 1: what a slow tier-1 costs, and what bounds it +# --------------------------------------------------------------------------- +class _Counted: + """A tier-1 client that counts calls and fails a chosen verb. + + `snippet_search` is the one that was down when this was written -- Asta's own + backend refusing a connection -- and the number that matters is not that it + failed but that it took 283 seconds to say so. + """ + + def __init__(self, fails=("snippet_search",), hits=1, seconds=0.0) -> None: + self.fails = set(fails) + self.hits = hits + self.seconds = seconds + self.calls: list[str] = [] + + def _call(self, verb: str, query: str, limit: int = 20): + import time + + self.calls.append(verb) + if self.seconds: + time.sleep(self.seconds) + if verb in self.fails: + raise UpstreamError(f"Asta's {verb} failed: ConnectionRefusedError", fix="retry") + return [ + {"id": f"s2:{query[:4]}-{i}", "paper_id": f"p{i}", "title": f"paper {i}", + "year": 2024, "snippet": "an excerpt", "source": "asta.paper", "external": {}} + for i in range(self.hits) + ] + + def snippet_search(self, query, limit=20): + return self._call("snippet_search", query, limit) + + def paper_search(self, query, limit=20): + return self._call("paper_search", query, limit) + + def neighbours(self, *_a, **_k): + return [] + + +def _expand_to(monkeypatch, queries): + """Stage 0, without Haiku: what matters here is how many queries stage 1 is + handed, because that is the multiplier on every tier-1 failure.""" + from core import haiku + + monkeypatch.setattr( + haiku, "expand", lambda *a, **k: {"queries": list(queries), "hyde": None} + ) + + +def test_a_dead_endpoint_is_dropped_after_one_failure_rather_than_per_query( + workspace, capsys, monkeypatch +): + """The run-killer. Stage 0 turns one question into six queries, and asking a + down endpoint once per query cost 283 seconds each time to learn the same + thing six times -- so the funnel was killed by its caller before the + endpoint that *was* working could return anything.""" + from core import http + from tools import paper_search + + client = _Counted(fails=("snippet_search",)) + monkeypatch.setattr(http, "PapersWithCode", lambda cfg: client) + _expand_to(monkeypatch, ["q1", "q2", "q3", "q4", "q5", "q6"]) + + code = paper_search.cli.run( + ["search", "efficient optimizers", "--no-triage", "--no-rerank", + "--no-local", "--no-citations", "--json"] + ) + payload = json.loads(capsys.readouterr().out) + + assert code == 0 + assert client.calls.count("snippet_search") == 1, "asked once, not once per query" + assert client.calls.count("paper_search") == 6, "the working endpoint is not punished" + assert payload["data"]["results"], "the run returns what the live endpoint found" + + +def test_what_was_dropped_is_written_into_the_trace(workspace, capsys, monkeypatch): + """A funnel that quietly searched one endpoint of two looks exactly like a + corpus that had little to say.""" + from core import http + from tools import paper_search + + monkeypatch.setattr(http, "PapersWithCode", lambda cfg: _Counted(fails=("snippet_search",))) + _expand_to(monkeypatch, ["q1", "q2"]) + paper_search.cli.run( + ["search", "efficient optimizers", "--no-triage", "--no-rerank", + "--no-local", "--no-citations", "--json"] + ) + capsys.readouterr() + + written = json.loads(next((paths.notes_dir() / "funnel").glob("*.json")).read_text("utf-8")) + assert written["stages"]["1_discovery"]["dropped"] == ["pwc.snippet_search"] + assert written["stages"]["1_discovery"]["queries_searched"] == 2 + assert any("dropped for the rest of this run" in w for w in written["warnings"]) + + +def test_stage_one_stops_when_its_wall_clock_is_spent_and_says_how_far_it_got( + workspace, capsys, monkeypatch +): + """`core/http.py` bounds one request; this bounds the run. Six queries over + two endpoints under a five-minute request deadline is a one-hour stage, and + the caller kills it long before it ends.""" + from core import http + from tools import paper_search + + config = paths.root() / "config" / "grad.toml" + config.parent.mkdir(parents=True, exist_ok=True) + config.write_text("[retrieval]\nstage1_budget_s = 0.05\n", encoding="utf-8") + + client = _Counted(fails=(), seconds=0.03) + monkeypatch.setattr(http, "PapersWithCode", lambda cfg: client) + _expand_to(monkeypatch, [f"q{i}" for i in range(20)]) + + code = paper_search.cli.run( + ["search", "efficient optimizers", "--no-triage", "--no-rerank", + "--no-local", "--no-citations", "--json"] + ) + payload = json.loads(capsys.readouterr().out) + + assert code == 0, "stopping early is not failing" + assert payload["data"]["results"], "what was retrieved is kept" + written = json.loads(next((paths.notes_dir() / "funnel").glob("*.json")).read_text("utf-8")) + searched = written["stages"]["1_discovery"]["queries_searched"] + assert 0 < searched < 20 + assert any("stage1_budget_s" in w for w in written["warnings"]) + + +def test_progress_reaches_stderr_so_a_pipe_can_see_the_run_moving( + workspace, capsys, monkeypatch +): + """Everything that runs this reads a pipe -- the tasks window streams the + tail, and the agent's own Bash gives up at 120s and backgrounds the command. + A run that prints nothing until its envelope is indistinguishable from a + hung one in both.""" + from core import http + from tools import paper_search + + monkeypatch.setattr(http, "PapersWithCode", lambda cfg: _Counted(fails=())) + _expand_to(monkeypatch, ["q1"]) + paper_search.cli.run( + ["search", "efficient optimizers", "--no-triage", "--no-rerank", + "--no-local", "--no-citations", "--json"] + ) + captured = capsys.readouterr() + assert "stage 1: pwc.paper_search" in captured.err + assert "stage 1" not in captured.out, "the --json contract stays one object on stdout" + + +def test_the_advice_matches_what_actually_failed(workspace): + """A live run failed with `ConnectionRefusedError` raised inside Asta's own + backend, and the fix said "discovery is rate limited" and pointed at a key. + No key affects that, and advice that cannot be followed is followed first.""" + from tools import paper_search + + limited = paper_search._tier1_fix(["asta"], ["asta.snippet_search: rate-limited"]) + assert "credential set asta_api_key" in limited + + refused = paper_search._tier1_fix( + ["pwc"], ["pwc.paper_search: ConnectionRefusedError inside the service"] + ) + assert "credential set" not in refused + assert "--local-only" in refused + + +def test_one_endpoint_being_down_does_not_hide_what_the_other_found( + workspace, capsys, monkeypatch +): + """The state of the service as this was written: `snippet_search` refuses + and `search_papers_by_relevance` works. A funnel that reports "every + retrieval call failed" in that situation is claiming the literature has + nothing on the question.""" + from core import http + from tools import paper_search + + monkeypatch.setattr(http, "PapersWithCode", lambda cfg: _Counted(fails=("snippet_search",), hits=3)) + _expand_to(monkeypatch, ["q1", "q2"]) + code = paper_search.cli.run( + ["search", "efficient optimizers", "--no-triage", "--no-rerank", + "--no-local", "--no-citations", "--json"] + ) + payload = json.loads(capsys.readouterr().out) + assert code == 0 + assert payload["ok"] is True + assert len(payload["data"]["results"]) == 6 + assert any("snippet_search" in w for w in payload["data"]["trace"]["warnings"]) diff --git a/tests/test_pwc.py b/tests/test_pwc.py new file mode 100644 index 0000000..7b5f4d3 --- /dev/null +++ b/tests/test_pwc.py @@ -0,0 +1,301 @@ +"""Papers with Code, and the arXiv batch that makes its rows readable. + +This replaced Asta as the funnel's default tier 1 for one reason, and it is a +measured one rather than a preference: Asta answers a search in ~121 seconds and +takes ~283 seconds to report that its own backend refused a connection, and +stage 0 multiplies that by six queries and two endpoints. Every caller gives up +first. This answers in one to two seconds. + +What it costs is the thing to keep honest, and most of what is asserted below is +about that: search rows carry **no abstract**, so `arxiv_abstracts` fills the +pool in one request, and `related` is a *dense neighbour* rather than a citation +edge and must not be presented as one. + +No network: the transport is faked, the same way `tests/test_asta.py` fakes it. +The shapes here are not guesses -- they are what the live API returned. +""" + +from __future__ import annotations + +import json + +import pytest + +from core import config as config_mod, http +from core.errors import UpstreamError + + +class FakeResponse: + def __init__(self, status_code=200, payload=None, text=""): + self.status_code = status_code + self._payload = payload + self.text = text or (json.dumps(payload) if payload is not None else "") + + def json(self): + if self._payload is None: + raise ValueError("not json") + return self._payload + + def raise_for_status(self): + if self.status_code >= 400: + raise RuntimeError(f"status {self.status_code}") + + +#: One row exactly as `papers/search` returned it live. Note what is *not* here. +LIVE_ROW = { + "id": "93318", + "title": "Towards Efficient Optimizer Design for LLM via Structured Fisher Approximation", + "arxiv_id": "2502.07752", + "source": "arxiv", + "authors": ["Wenbo Gong", "Meyer Scetbon"], + "published": "2025-02-11", + "url_abs": "https://arxiv.org/abs/2502.07752", + "citation_count": 6, +} + + +@pytest.fixture +def transport(monkeypatch): + """A fake catalogue. `queue` answers each GET in order.""" + gets: list[dict] = [] + queue: list[FakeResponse] = [] + + class FakeHttpx: + @staticmethod + def get(url, params=None, headers=None, timeout=None): + gets.append({"url": url, "params": params or {}, "headers": headers}) + return queue.pop(0) if queue else FakeResponse(payload={"results": []}) + + monkeypatch.setattr(http, "_httpx", lambda: FakeHttpx) + return gets, queue + + +def client(): + return http.PapersWithCode(config_mod.load(reload=True)) + + +# --------------------------------------------------------------------------- +# the two rankings +# --------------------------------------------------------------------------- +def test_the_two_verbs_are_two_genuinely_different_rankings(workspace, transport): + """Lexical and dense, which is the pair `corpus.rrf` exists to fuse -- so the + two calls the funnel already makes per query map onto them without the + funnel knowing anything changed.""" + gets, queue = transport + queue.append(FakeResponse(payload={"results": [LIVE_ROW]})) + queue.append(FakeResponse(payload={"results": [LIVE_ROW]})) + + client().snippet_search("efficient optimizers") + client().paper_search("efficient optimizers") + assert gets[0]["params"]["mode"] == "semantic" + assert gets[1]["params"]["mode"] == "keyword" + + +def test_a_row_arrives_in_the_vocabulary_the_funnel_already_fuses(workspace, transport): + _, queue = transport + queue.append(FakeResponse(payload={"results": [LIVE_ROW]})) + row = client().paper_search("efficient optimizers")[0] + + shared = {"id", "paper_id", "title", "year", "snippet", "abstract", "source", "external"} + assert shared <= set(row) + assert row["id"] == "pwc:2502.07752" + assert row["paper_id"] == "2502.07752" + assert row["year"] == "2025" + assert row["external"]["ArXiv"] == "2502.07752" + # The honest part: the catalogue's search does not return text. + assert row["snippet"] == "" + assert row["abstract"] == "" + + +def test_this_catalogue_does_not_share_the_semantic_scholar_namespace(workspace, transport): + """Asta and S2 share `s2:` because they are one index with one set of ids. + This is a different catalogue with its own numbering, and giving it the same + prefix would fuse two unrelated papers whose ids happened to collide.""" + _, queue = transport + queue.append(FakeResponse(payload={"results": [{"id": "991", "title": "Attention"}]})) + assert client().paper_search("attention")[0]["id"] == "pwc:991" + + +def test_a_page_larger_than_the_api_allows_is_clamped(workspace, transport): + gets, queue = transport + queue.append(FakeResponse(payload={"results": []})) + client().paper_search("attention", limit=500) + assert gets[0]["params"]["page_size"] == http.PapersWithCode.MAX_PAGE + + +# --------------------------------------------------------------------------- +# expansion, and what it is not +# --------------------------------------------------------------------------- +def test_related_work_is_reported_as_a_dense_neighbour_not_a_citation(workspace, transport): + """The rows carry `provenance: "dense"` and a similarity score: they are + nearest neighbours in an embedding space, not papers that cite the seed. + A ledger entry's basis is the thing that must not be overstated.""" + _, queue = transport + queue.append(FakeResponse(payload=[{**LIVE_ROW, "provenance": "dense", "similarity": 0.86}])) + row = client().neighbours("1706.03762")[0] + assert row["source"] == "pwc.related" + assert "citation" not in row["source"] + + +def test_the_backward_direction_is_refused_rather_than_faked(workspace, transport): + gets, _ = transport + assert client().neighbours("1706.03762", direction="references") == [] + assert gets == [], "no request should have been made at all" + + +def test_related_answers_with_a_bare_list_and_search_with_an_envelope(workspace, transport): + """Both are what the live API returns, and both have to be read.""" + _, queue = transport + queue.append(FakeResponse(payload=[LIVE_ROW])) + assert len(client().neighbours("1706.03762")) == 1 + + +def test_an_unrecognised_shape_raises_rather_than_returning_nothing(workspace, transport): + """`ok: true` with no results reads as "the literature has nothing on this", + which is the one conclusion a schema change must not manufacture.""" + _, queue = transport + queue.append(FakeResponse(payload={"papers_maybe": []})) + with pytest.raises(UpstreamError) as exc: + client().paper_search("attention") + assert "pwc-cli" in (exc.value.fix or "") + + +def test_rate_limiting_says_there_is_no_key_to_add(workspace, transport): + _, queue = transport + queue.append(FakeResponse(status_code=429, text="slow down")) + with pytest.raises(UpstreamError) as exc: + client().paper_search("attention") + assert "anonymous" in (exc.value.fix or "") + + +def test_a_moved_endpoint_names_the_config_key_that_moves_with_it(workspace, transport): + _, queue = transport + queue.append(FakeResponse(status_code=404, text="gone")) + with pytest.raises(UpstreamError) as exc: + client().paper_search("attention") + assert "pwc_base" in (exc.value.fix or "") + + +# --------------------------------------------------------------------------- +# the abstracts +# --------------------------------------------------------------------------- +ATOM = """ + + + http://arxiv.org/abs/2211.09760v1 + VeLO + While deep learning models have replaced + hand-designed features, these models are still trained. + + + http://arxiv.org/abs/1706.03762v7 + Attention Is All You Need + The dominant sequence transduction models. + +""" + + +def test_a_hundred_abstracts_cost_one_request(workspace, transport): + """The whole reason to use `id_list`: fetching them one at a time would cost + a request per candidate and undo the reason the corpus was changed.""" + gets, queue = transport + queue.append(FakeResponse(text=ATOM)) + out = http.arxiv_abstracts( + ["2211.09760", "1706.03762"], cfg=config_mod.load(reload=True) + ) + assert len(gets) == 1 + assert gets[0]["params"]["id_list"] == "2211.09760,1706.03762" + assert out["2211.09760"].startswith("While deep learning models") + # Flattened: the feed wraps them, and a reranker reading newlines as + # structure would be reading the feed's formatting. + assert "\n" not in out["2211.09760"] + + +def test_the_version_suffix_is_stripped_so_the_lookup_matches(workspace, transport): + """The feed answers `.../abs/1706.03762v7` and the caller asked for + `1706.03762`, which is what its candidates are keyed by.""" + _, queue = transport + queue.append(FakeResponse(text=ATOM)) + out = http.arxiv_abstracts(["1706.03762"], cfg=config_mod.load(reload=True)) + assert "1706.03762" in out + + +def test_a_failed_fetch_degrades_the_ranking_rather_than_the_run(workspace, transport): + """A candidate with no abstract ranks on its title, which is what would have + happened without this. What it must not do is raise.""" + _, queue = transport + queue.append(FakeResponse(status_code=503, text="down")) + assert http.arxiv_abstracts(["1706.03762"], cfg=config_mod.load(reload=True)) == {} + + +def test_nothing_to_fetch_is_not_a_request(workspace, transport): + gets, _ = transport + assert http.arxiv_abstracts([], cfg=config_mod.load(reload=True)) == {} + assert http.arxiv_abstracts(["", None], cfg=config_mod.load(reload=True)) == {} + assert gets == [] + + +# --------------------------------------------------------------------------- +# the funnel, over the whole of it +# --------------------------------------------------------------------------- +def test_the_pool_is_enriched_before_it_is_reranked(workspace, capsys, monkeypatch): + """Stage 2 reranks on `title + snippet-or-abstract` and stage 3 triages on + the same, so a pool of bare titles is a measurably worse funnel -- and a + funnel silently reranking titles looks exactly like one reranking + abstracts.""" + from tools import paper_search + + class Catalogue: + def snippet_search(self, query, limit=20): + return [] + + def paper_search(self, query, limit=20): + return [{ + "id": "pwc:2502.07752", "paper_id": "2502.07752", "title": "A paper", + "year": "2025", "snippet": "", "abstract": "", "section": "", + "source": "pwc.keyword", "external": {"ArXiv": "2502.07752"}, + }] + + def neighbours(self, *_a, **_k): + return [] + + monkeypatch.setattr(http, "PapersWithCode", lambda cfg: Catalogue()) + monkeypatch.setattr( + http, "arxiv_abstracts", lambda ids, **_k: {"2502.07752": "the real abstract"} + ) + paper_search.cli.run( + ["search", "optimizers", "--no-expand", "--no-rerank", "--no-triage", + "--no-local", "--no-citations", "--full", "--json"] + ) + payload = json.loads(capsys.readouterr().out) + assert payload["data"]["results"][0]["text"] == "the real abstract" + assert payload["data"]["trace"]["stages"]["1_abstracts"] == {"wanted": 1, "found": 1} + + +def test_candidates_left_without_an_abstract_are_reported(workspace, capsys, monkeypatch): + from tools import paper_search + + class Catalogue: + def snippet_search(self, query, limit=20): + return [] + + def paper_search(self, query, limit=20): + return [{ + "id": "pwc:1", "paper_id": "9999.99999", "title": "A paper", "year": "2025", + "snippet": "", "abstract": "", "section": "", "source": "pwc.keyword", + "external": {"ArXiv": "9999.99999"}, + }] + + def neighbours(self, *_a, **_k): + return [] + + monkeypatch.setattr(http, "PapersWithCode", lambda cfg: Catalogue()) + monkeypatch.setattr(http, "arxiv_abstracts", lambda ids, **_k: {}) + paper_search.cli.run( + ["search", "optimizers", "--no-expand", "--no-rerank", "--no-triage", + "--no-local", "--no-citations", "--json"] + ) + payload = json.loads(capsys.readouterr().out) + assert any( + "title alone" in w for w in payload["data"]["trace"]["warnings"] + ), "a funnel reranking titles must not look like one reranking abstracts" diff --git a/tests/test_streaming.py b/tests/test_streaming.py index ce51260..61d5f06 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -451,3 +451,128 @@ def test_the_options_actually_ask_for_partial_messages(monkeypatch): options = agent.build_options(config_mod.load()) assert isinstance(options, sdk.ClaudeAgentOptions) assert options.include_partial_messages is True + + +# --------------------------------------------------------------------------- +# the reasoning (`TurnStream`, the third kind of block) +# --------------------------------------------------------------------------- +def thinking(text: str) -> FakeStreamEvent: + return FakeStreamEvent( + {"type": "content_block_delta", "delta": {"type": "thinking_delta", "thinking": text}} + ) + + +def test_the_reasoning_is_a_block_of_its_own_and_not_part_of_the_answer(): + """The chat window's statusline switches this on and off, so it has to be + separable. `text` is what the agent *said*; the working is a different + claim and the transcript records it as one.""" + stream = agent.TurnStream() + stream.feed(thinking("the ceiling is the binding constraint")) + stream.feed(delta("Raise the ceiling.")) + stream.feed(FakeMessage([ + ThinkingBlock("the ceiling is the binding constraint"), + FakeBlock("Raise the ceiling."), + ])) + assert kinds(stream) == ["thinking", "text"] + assert stream.text == "Raise the ceiling." + assert stream.thinking == "the ceiling is the binding constraint" + + +def test_reasoning_is_not_printed_by_the_command_line(): + """`feed` returns what a CLI prints. The reasoning is a block a UI may draw, + not something the terminal session starts emitting.""" + stream = agent.TurnStream() + printed = stream.feed(thinking("hmm")) + stream.feed(delta("Yes.")) + assert printed == "Yes." + + +def test_a_finished_thinking_block_does_not_repeat_the_deltas_that_built_it(): + """The same trap `TextStream` exists for, one channel over: the SDK sends + `thinking_delta` events *and* the `ThinkingBlock` containing all of them.""" + stream = agent.TurnStream() + stream.feed(thinking("first ")) + stream.feed(thinking("second")) + stream.feed(FakeMessage([ThinkingBlock("first second")])) + assert stream.thinking == "first second" + assert kinds(stream) == ["thinking"] + + +def test_reasoning_resumes_below_a_call_rather_than_above_it(): + """Interleaved thinking: the order is the information here too. Reasoning + that happened *after* a command must not be appended to the block that was + open before it ran.""" + stream = agent.TurnStream() + stream.feed(FakeMessage([ThinkingBlock("check the ledger first")])) + stream.feed(FakeMessage([ToolUse("tu_1", "Bash", {"command": "ls"})])) + stream.feed(FakeMessage([ToolResult("tu_1", "one")])) + stream.feed(FakeMessage([ThinkingBlock("one entry, so the claim holds")])) + assert kinds(stream) == ["thinking", "tool", "thinking"] + assert [b["text"] for b in stream.blocks if b["kind"] == "thinking"] == [ + "check the ledger first", + "one entry, so the claim holds", + ] + + +def test_a_turn_that_only_reasoned_still_settles_as_something(): + """`Session.ask` keeps a turn when it produced blocks. Reasoning is blocks, + so an interrupted turn that had only got as far as thinking is not recorded + as a prompt that went unanswered.""" + stream = agent.TurnStream() + stream.feed(thinking("still working out what to run")) + assert stream.text == "" + assert stream.blocks and stream.blocks[0]["kind"] == "thinking" + + +def test_a_call_is_stamped_so_the_tasks_window_can_age_it(): + """Wall clock, not monotonic: it goes into the session file and is read back + in another process, where a monotonic reading means nothing.""" + import time + + stream = agent.TurnStream() + stream.feed(FakeMessage([ToolUse("tu_1", "Bash", {"command": "sleep 60"})])) + started = stream.blocks[0]["started"] + assert isinstance(started, float) + assert abs(started - time.time()) < 60 + + +def test_the_reasoning_is_asked_for_as_text_rather_than_assumed(monkeypatch): + """Capturing thinking blocks is not enough to have any: Opus 4.7+ defaults + `display` to "omitted" and sends them with a signature and no text. The chat + window's reasoning switch had nothing to reveal no matter how correctly the + stream was read -- one flag away from the feature, and indistinguishable + from a toggle that does nothing.""" + sdk = pytest.importorskip("claude_agent_sdk", reason="the SDK is not installed") + from core import config as config_mod + + monkeypatch.setattr(agent, "system_prompt", lambda: "prompt") + options = agent.build_options(config_mod.load(reload=True)) + assert options.thinking == {"type": "adaptive", "display": "summarized"} + + +def test_reasoning_can_be_turned_off_from_the_config(workspace, monkeypatch): + sdk = pytest.importorskip("claude_agent_sdk", reason="the SDK is not installed") + from core import config as config_mod, paths + + config = paths.root() / "config" / "grad.toml" + config.parent.mkdir(parents=True, exist_ok=True) + config.write_text('[agent]\nreasoning = "omitted"\n', encoding="utf-8") + monkeypatch.setattr(agent, "system_prompt", lambda: "prompt") + options = agent.build_options(config_mod.load(reload=True)) + assert options.thinking == {"type": "adaptive", "display": "omitted"} + + +def test_an_sdk_without_the_option_still_builds_a_session(monkeypatch): + """Feature-detected rather than assumed, for the same reason the deny probe + exists: this option is newer than the permission mode, and the SDK's shape + has changed between releases.""" + import dataclasses + + from core import config as config_mod + + class OldOptions: + pass + + class OldSdk: + ClaudeAgentOptions = dataclasses.make_dataclass("ClaudeAgentOptions", ["model"]) + + assert agent.thinking_option(config_mod.load(reload=True), OldSdk) == {} diff --git a/tests/test_ui_shell.py b/tests/test_ui_shell.py index 936693f..d9198aa 100644 --- a/tests/test_ui_shell.py +++ b/tests/test_ui_shell.py @@ -427,14 +427,66 @@ def test_a_new_turn_clears_the_tail_rather_than_stacking_onto_it(rendered): def test_the_status_line_names_the_call_in_flight(rendered): """A spinner says something is happening; naming the command says a - 40-minute job is running and which one.""" + 40-minute job is running and which one. + + The fallbacks changed with the statusline: `running …` was the only thing it + could say when no call was in flight, which covered "reasoning", "writing" + and "nothing has come back yet" with one word. Each is now named, because + each is a different answer to "why is nothing on screen". + """ from ui.windows.chat import _activity - assert _activity([]) == "running …" + assert _activity([]) == "waiting for the model" assert _activity([ {"kind": "tool", "name": "Bash", "title": "python -m tools.jobs run", "status": "running"}, ]) == "running Bash python -m tools.jobs run" - assert _activity([{"kind": "tool", "name": "Bash", "title": "ls", "status": "ok"}]) == "running …" + # A finished call is not an activity; what is happening is whatever came + # after it, and the reasoning is as specific as that gets. + assert _activity([ + {"kind": "tool", "name": "Bash", "title": "ls", "status": "ok"}, + {"kind": "thinking", "text": "one entry, so the claim holds"}, + ]) == "thinking" + assert _activity([ + {"kind": "thinking", "text": "working it out"}, + {"kind": "text", "text": "The answer is"}, + ]) == "writing" + + +def test_the_statusline_switches_the_reasoning_without_redrawing_the_transcript(rendered): + """A toggle that rebuilt the transcript would take its scroll position with + it, which is the same reason the poll never touches this window. So the + blocks are always in the DOM and a class decides whether they are painted.""" + client, space = rendered(["chat"]) + with client: + assert space.show_reasoning is False + roots = [ + e for e in client.elements.values() if "grad-chat" in getattr(e, "classes", []) + ] + assert roots, "the chat root carries the class the switch writes" + assert "reasoning-on" not in roots[0].classes + + bars = [ + e for e in client.elements.values() + if "grad-statusline" in getattr(e, "classes", []) + ] + assert len(bars) == 1 + + assert space.toggle_reasoning() is True + roots[0].classes(add="reasoning-on") + assert "reasoning-on" in roots[0].classes + + +def test_the_session_picker_is_the_workspaces_own_menu(rendered): + """The last Quasar control in the workspace. A `select` has one string per + option, so "another window has this open" arrived as a ` · ` fragment glued + onto the title and looked exactly like the rows that can be opened.""" + client, space = rendered(["chat"]) + with client: + selects = [ + e for e in client.elements.values() if type(e).__name__ == "Select" + ] + assert selects == [] + assert "grad-session-btn" in html_of(client) def test_the_focused_window_is_marked(rendered): diff --git a/tests/test_ui_tasks.py b/tests/test_ui_tasks.py index 125b792..4c2be10 100644 --- a/tests/test_ui_tasks.py +++ b/tests/test_ui_tasks.py @@ -331,3 +331,106 @@ async def test_run_tool_reports_a_command_that_printed_no_envelope(workspace): payload = await tasks_mod.run_tool(name) assert payload["ok"] is False assert "something went wrong" in payload["error"]["message"] + + +# --------------------------------------------------------------------------- +# the agent's own calls +# --------------------------------------------------------------------------- +class FakeSession: + """A `ui.app.Session` as far as the tasks model is concerned: the turn in + flight, and the turns that settled.""" + + def __init__(self, blocks=None, settled=None) -> None: + self.blocks = blocks or [] + self.settled = settled or [] + + +def call(cid, name="Bash", title="ls", status="running", result="", started=None): + import time as _t + + block = { + "kind": "tool", "id": cid, "name": name, "title": title, + "status": status, "result": result, + } + block["started"] = _t.time() if started is None else started + return block + + +def test_the_turn_in_flight_is_what_the_tasks_window_shows_first(): + """Every capability in this project is reached by a Bash into `tools/`, so + the agent's calls are the other half of "what is running on this machine". + Until this they were visible only in the transcript, which is the wrong + place to look once the conversation has scrolled on.""" + from ui import models + + session = FakeSession( + blocks=[{"kind": "text", "text": "checking"}, call("tu_2", title="pytest -q")], + settled=[{"role": "assistant", "blocks": [call("tu_1", status="ok", result="one\ntwo")]}], + ) + rows = models.agent_calls_model(session) + assert [r["id"] for r in rows] == ["tu_2", "tu_1"] + assert rows[0]["state"] == "running" + assert rows[1]["state"] == "ok" + + +def test_a_call_left_running_by_a_settled_turn_is_not_reported_as_running(): + """The turn died or was interrupted mid-call. Whatever it started is not + this app's to know about, and saying "running" of something nothing is + waiting for is the same lie as a tail that silently forgets.""" + from ui import models + + session = FakeSession(settled=[{"role": "assistant", "blocks": [call("tu_1")]}]) + row = models.agent_calls_model(session)[0] + assert row["state"] == "unfinished" + assert row["running"] is False + assert row["elapsed"] == "" + + +def test_only_a_live_call_is_given_a_clock(): + from ui import models + + session = FakeSession(blocks=[call("tu_1", started=None)]) + assert models.agent_calls_model(session)[0]["elapsed"] != "" + + +def test_a_call_from_a_transcript_written_before_calls_were_stamped_still_lists(): + """The session file outlives the version that wrote it. The row is worth + showing; the clock is the part that is not known.""" + from ui import models + + block = call("tu_1") + del block["started"] + assert models.agent_calls_model(FakeSession(blocks=[block]))[0]["elapsed"] == "" + + +def test_the_call_list_is_bounded_so_the_window_is_not_a_second_transcript(): + from ui import models + + settled = [ + {"role": "assistant", "blocks": [call(f"tu_{i}", status="ok")]} + for i in range(models.AGENT_CALLS * 2) + ] + assert len(models.agent_calls_model(FakeSession(settled=settled))) == models.AGENT_CALLS + + +def test_the_two_lists_are_counted_apart(): + """A task is a process this app started and can stop; a call is one the + agent made and only the agent can stop. Merging them would imply a STOP + button that does not exist.""" + from ui import models, tasks + + tasks.start("a wiki rebuild", "tools.wiki", "map") + model = models.tasks_model(agent=models.agent_calls_model(FakeSession(blocks=[call("tu_1")]))) + assert model["running"] == 1 + assert model["agent_running"] == 1 + assert [r["id"] for r in model["rows"]] != [c["id"] for c in model["agent"]] + + +def test_the_tasks_model_without_a_session_is_unchanged(): + """`tasks_model` is called from the poll and from tests with no session at + all; the agent half is additive.""" + from ui import models + + model = models.tasks_model() + assert model["agent"] == [] + assert model["agent_running"] == 0 diff --git a/tests/test_ui_turns.py b/tests/test_ui_turns.py new file mode 100644 index 0000000..afb93c9 --- /dev/null +++ b/tests/test_ui_turns.py @@ -0,0 +1,252 @@ +"""Interrupting a turn, and the turn after it (`ui.app.Session`). + +The bug this suite exists for, as reported: interrupt the agent, submit another +prompt, and the answer is invisible until you interrupt *again*. It has one +visible symptom and three separate causes, all of which end in the same place -- +`Session.busy` is still True, so the composer's guard silently refuses the next +prompt and nothing on screen says why: + + * the SDK refused the interrupt, and every exception was swallowed; + * the interrupt was accepted and the turn did not end; + * the interrupt landed *late*, on the turn issued after the one it was aimed + at, killing it the moment it started. + +None of the three needs the SDK to reproduce -- they are all about what this +class does with a client that behaves in a particular way -- so the client here +is a fake that can be told to behave in each of them. +""" + +from __future__ import annotations + +import asyncio + +import pytest + +app = pytest.importorskip("ui.app", reason="the ui extra is not installed") + +pytestmark = pytest.mark.asyncio + + +class FakeClient: + """A `ClaudeSDKClient` shaped just enough to drive one turn. + + `receive_response` yields nothing and waits until `finish` is set, which is + what an in-flight turn is; `interrupt` sets it, which is what a working + interrupt does. The variations are what the tests are about. + """ + + def __init__(self, *, refuses: bool = False, ignores: bool = False) -> None: + self.refuses = refuses + self.ignores = ignores + self.finish = asyncio.Event() + self.prompts: list[str] = [] + self.interrupts = 0 + self.closed = False + + async def __aenter__(self): + return self + + async def __aexit__(self, *_): + self.closed = True + # Whatever was waiting for a message is not going to get one. + self.finish.set() + return False + + async def query(self, prompt: str) -> None: + self.prompts.append(prompt) + self.finish.clear() + + async def receive_response(self): + await self.finish.wait() + return + yield # pragma: no cover - makes this an async generator + + async def interrupt(self) -> None: + self.interrupts += 1 + if self.refuses: + raise RuntimeError("not connected") + if not self.ignores: + self.finish.set() + + +@pytest.fixture +def session(monkeypatch): + """A `Session` whose client is a fake and whose notices are collected.""" + notices: list[str] = [] + made: list[FakeClient] = [] + kind: dict[str, dict] = {"kwargs": {}} + + made_session = app.Session("turns") + made_session.notify = notices.append + + async def start() -> None: + if made_session.client is None: + made_session.client = FakeClient(**kind["kwargs"]) + made.append(made_session.client) + + monkeypatch.setattr(made_session, "start", start) + made_session.notices = notices # type: ignore[attr-defined] + made_session.clients = made # type: ignore[attr-defined] + made_session.client_kind = kind # type: ignore[attr-defined] + return made_session + + +async def settled(_record) -> None: + return None + + +async def run_turn(session, prompt: str = "hello"): + """Start a turn and wait until it is genuinely in flight.""" + task = asyncio.create_task(session.ask(prompt, settled)) + for _ in range(200): + await asyncio.sleep(0) + if session.client is not None and session.client.prompts: + return task + raise AssertionError("the turn never reached the client") + + +# --------------------------------------------------------------------------- +# the three causes +# --------------------------------------------------------------------------- +async def test_an_interrupt_the_sdk_refuses_is_reported_and_still_ends_the_turn(session, monkeypatch): + """Every exception used to be swallowed here, which is precisely how + pressing STOP twice became the way to stop a turn: the first press failed in + silence and the second one happened to land. + + A refused interrupt is also the case that most needs the escalation: nothing + asked the turn to stop, so nothing will end it but taking the client down. + """ + monkeypatch.setattr(app, "INTERRUPT_GRACE_S", 0.05) + session.client_kind["kwargs"] = {"refuses": True} + task = await run_turn(session) + + assert session.interrupt() == "interrupting the turn…" + await asyncio.wait_for(task, timeout=5) + + assert session.busy is False + assert any("refused the interrupt" in n for n in session.notices) + + +async def test_a_turn_that_ignores_the_interrupt_is_taken_down_anyway(session, monkeypatch): + """`ui/tasks.py:cancel`'s escalation, applied to the one control that did + not have it. A session that stays busy forever refuses every prompt after + it, and says nothing about either.""" + monkeypatch.setattr(app, "INTERRUPT_GRACE_S", 0.05) + session.client_kind["kwargs"] = {"ignores": True} + task = await run_turn(session) + client = session.client + + session.interrupt() + await asyncio.wait_for(task, timeout=5) + + assert client.interrupts == 1, "the tool's own stop is asked for first" + assert client.closed is True + assert session.busy is False + assert any("did not stop" in n for n in session.notices) + + +async def test_the_next_turn_waits_for_a_pending_interrupt(session, monkeypatch): + """The late interrupt. Fire-and-forget, it outlives the turn it was aimed at + and lands on the one after it -- which then produces nothing and explains + nothing, and is the shape of the bug as reported.""" + monkeypatch.setattr(app, "INTERRUPT_GRACE_S", 0.05) + session.client_kind["kwargs"] = {"ignores": True} + first = await run_turn(session, "one") + + session.interrupt() + await asyncio.wait_for(first, timeout=5) + + second = await run_turn(session, "two") + # The interrupt is finished before the second turn is issued, so it cannot + # be the thing that ends it. + assert session._stopping is None or session._stopping.done() # noqa: SLF001 + assert session.client.interrupts == 0 + session.client.finish.set() + await asyncio.wait_for(second, timeout=5) + assert session.client.prompts == ["two"] + + +# --------------------------------------------------------------------------- +# what the control says +# --------------------------------------------------------------------------- +async def test_interrupting_nothing_says_so_rather_than_pretending(session): + assert session.interrupt() == "nothing is running" + + +async def test_a_second_press_does_not_stack_a_second_interrupt(session, monkeypatch): + monkeypatch.setattr(app, "INTERRUPT_GRACE_S", 0.05) + session.client_kind["kwargs"] = {"ignores": True} + task = await run_turn(session) + + session.interrupt() + assert session.interrupt() == "already stopping — the turn is being taken down" + await asyncio.wait_for(task, timeout=5) + assert session.clients[0].interrupts == 1 + + +# --------------------------------------------------------------------------- +# what survives it +# --------------------------------------------------------------------------- +async def test_the_client_is_rebuilt_after_an_interrupt_and_resumes_the_conversation(session): + """A fresh client cannot hold a message belonging to the turn that was + stopped -- which is the other way the next turn ended instantly with nothing + drawn. `resume` is what keeps that from costing the conversation, so the id + has to have been recorded by then.""" + task = await run_turn(session) + session.sdk_session_id = "sdk-1" + + session.interrupt() + await asyncio.wait_for(task, timeout=5) + assert session.client is None + + second = await run_turn(session, "again") + assert len(session.clients) == 2, "a new client, not the interrupted one" + session.client.finish.set() + await asyncio.wait_for(second, timeout=5) + assert session.sdk_session_id == "sdk-1" + + +async def test_the_sdk_session_id_is_recorded_from_a_turn_that_never_finished(session, monkeypatch): + """`drive_turn`'s return value is not reached on the interrupt path, and + that id is what the rebuilt client resumes from -- so reading it off the + return value lost the conversation on exactly the turns that end in a + rebuild.""" + import agent as agent_mod + + class Message: + session_id = "sdk-7" + + async def drive_turn(client, prompt, stream, *, on_session_id=None, **_): + on_session_id(Message.session_id) + raise RuntimeError("interrupted") + + monkeypatch.setattr(agent_mod, "drive_turn", drive_turn) + await session.ask("hello", settled) + assert session.sdk_session_id == "sdk-7" + + +async def test_a_partly_streamed_turn_keeps_what_it_streamed(session, monkeypatch): + """An interrupted turn is more legible with its half than without it, and + the transcript must not claim the prompt went unanswered.""" + import agent as agent_mod + + async def drive_turn(client, prompt, stream, **_): + stream.feed(_FakeAssistant("half an answer")) + raise RuntimeError("interrupted") + + monkeypatch.setattr(agent_mod, "drive_turn", drive_turn) + await session.ask("hello", settled) + + assert session.settled[0] == {"role": "user", "text": "hello"} + assert "half an answer" in session.settled[1]["text"] + assert "the session failed" in session.settled[1]["text"] + + +class _FakeAssistant: + def __init__(self, text: str) -> None: + self.content = [_FakeTextBlock(text)] + + +class _FakeTextBlock: + def __init__(self, text: str) -> None: + self.text = text diff --git a/tools/paper_search.py b/tools/paper_search.py index 68145b1..414f0bd 100644 --- a/tools/paper_search.py +++ b/tools/paper_search.py @@ -17,6 +17,8 @@ import argparse import json import re +import sys +import time from typing import Any from core import config as config_mod, corpus, credentials, haiku, http, paths, quota_log @@ -40,28 +42,82 @@ #: `tier1` value -> the clients it selects, in the order they are queried. -TIER1_SOURCES = ("asta", "s2", "both", "none") +#: `both` predates `pwc` and still means what it meant: the two Semantic Scholar +#: doors, for comparing them. +TIER1_SOURCES = ("pwc", "asta", "s2", "both", "all", "none") + + +class _Budget: + """A wall clock over the whole of stage 1. + + `core/http.py` bounds one *request*; this bounds the run. They are different + numbers because stage 0 turns one question into six queries and each is put + to two endpoints, so a per-request deadline of five minutes is a + one-hour stage. What made that concrete: a live `snippet_search` took 283 + seconds to report that Asta's own backend had refused a connection, and the + funnel was killed by its caller long before the endpoints that *were* + working could contribute anything. + + Stopping early is not the same as failing. What has been retrieved is kept, + what was skipped is written into the trace, and the caller gets results -- + which is the whole difference between a slow funnel and a broken one. + """ + + def __init__(self, limit: float) -> None: + self.limit = max(0.0, limit) + self._started = time.monotonic() + + @property + def elapsed(self) -> float: + return time.monotonic() - self._started + + @property + def spent(self) -> bool: + return bool(self.limit) and self.elapsed >= self.limit + + +def _progress(message: str) -> None: + """A line of progress on stderr, where a `--json` contract cannot see it. + + The funnel prints nothing until its envelope, which is minutes later, and + everything that runs it reads a pipe: the tasks window streams the tail, and + the agent's own Bash call gives up at 120 seconds and moves the command to + the background. A run with no output is indistinguishable from a hung one in + all three places, and that is how a slow stage got read as a broken tool. + """ + print(message, file=sys.stderr, flush=True) def tier1_clients(cfg: Any, override: str | None = None) -> list[tuple[str, Any]]: """The discovery clients for this run, named so a trace can say which spoke. - Both reach the same Semantic Scholar corpus and both answer in the same - vocabulary (`core/http.py:_row`), so a candidate found by either fuses to - one entry. What differs is whether the door opens: S2's own API no longer - issues keys to free-domain addresses, so a personal account falls back to - the shared anonymous pool. + All three answer in the same vocabulary (`core/http.py:_row`, `_pwc_row`), + so the funnel does not know which tier a candidate came from and does not + have to. What differs is whether the door opens, and how fast: + + * **pwc** -- Papers with Code, as revived by Hugging Face. Anonymous, one to + two seconds, and two genuinely different rankings (lexical and dense) that + RRF was built to fuse. The default, because it is the one that answers. + * **asta** -- the same Semantic Scholar corpus as `s2`, through a door that + opens without an institutional address, and the only one of the three with + real full-text snippets. Measured live at 121s for a search and 283s to + report a backend failure, which is why it is no longer the default. + * **s2** -- the REST API directly. Its own keys are only issued to + institutional addresses, so a personal account falls back to the shared + anonymous pool, which is near-permanently rate limited. """ - chosen = str(override or cfg.get("retrieval", "tier1", "asta")).lower() + chosen = str(override or cfg.get("retrieval", "tier1", "pwc")).lower() if chosen not in TIER1_SOURCES: raise UsageError( f"unknown tier-1 source {chosen!r}", fix=f"one of: {', '.join(TIER1_SOURCES)}", ) out: list[tuple[str, Any]] = [] - if chosen in ("asta", "both"): + if chosen in ("pwc", "all"): + out.append(("pwc", http.PapersWithCode(cfg))) + if chosen in ("asta", "both", "all"): out.append(("asta", http.Asta(cfg))) - if chosen in ("s2", "both"): + if chosen in ("s2", "both", "all"): out.append(("s2", http.SemanticScholar(cfg))) return out @@ -120,19 +176,60 @@ def cmd_search(args: argparse.Namespace) -> dict[str, Any]: if tier1: per_query = max(5, ceiling // max(1, len(queries) * 2 * len(tier1))) + budget = _Budget(float(cfg.get("retrieval", "stage1_budget_s", 300))) + #: `(source, verb)` pairs that have already failed once. See `_Budget`. + dropped: set[tuple[str, str]] = set() + searched = 0 for query in queries: + if budget.spent: + break + searched += 1 for name, client in tier1: for verb in ("snippet_search", "paper_search"): + if (name, verb) in dropped or budget.spent: + continue + _progress(f"stage 1: {name}.{verb} q{searched}/{len(queries)}") try: hits = getattr(client, verb)(query, limit=per_query) except GradError as exc: + # Dropped for the rest of the run, not merely skipped for + # this query. A tier-1 endpoint that is down is down for + # every query, and finding that out costs a *request* -- + # a live snippet_search took 283 seconds to report that + # Asta's own backend had refused a connection, so trying + # it once per expanded query spent 28 minutes learning + # the same thing six times and got the run killed before + # the stages that were working could return anything. + dropped.add((name, verb)) trace.setdefault("warnings", []).append(f"{name}.{verb}: {exc}") upstream_failures.append(f"{name}.{verb}: {exc}") + _progress(f"stage 1: {name}.{verb} failed; not retried this run") continue rankings.append(hits) + _progress(f"stage 1: {name}.{verb} returned {len(hits)}") for hit in hits: candidates.setdefault(hit["id"], hit) - if not args.no_citations: + # Both caps are reported rather than left to be inferred from a thin + # result set: a funnel that quietly searched two of six queries looks + # exactly like a corpus that had little to say. + if dropped: + trace.setdefault("warnings", []).append( + "dropped for the rest of this run after one failure each: " + + ", ".join(sorted(f"{n}.{v}" for n, v in dropped)) + ) + if budget.spent: + trace.setdefault("warnings", []).append( + f"tier-1 discovery stopped after {budget.limit:.0f}s having searched " + f"{searched} of {len(queries)} queries — raise [retrieval] " + "stage1_budget_s, or --no-expand to search one query instead of six" + ) + trace["stages"]["1_discovery"] = { + "queries_searched": searched, + "queries": len(queries), + "dropped": sorted(f"{n}.{v}" for n, v in dropped), + "seconds": round(budget.elapsed, 1), + } + if not args.no_citations and not budget.spent: seeds = [c for c in list(candidates.values())[:5] if c.get("paper_id")] for seed in seeds: for name, client in tier1: @@ -156,6 +253,7 @@ def cmd_search(args: argparse.Namespace) -> dict[str, Any]: fused = corpus.rrf(rankings, k=int(cfg.get("retrieval", "rrf_k", 60))) pool = [candidates[f["id"]] | {"rrf": f["rrf"]} for f in fused if f["id"] in candidates][:ceiling] + _fill_abstracts(pool, cfg, trace) quota_log.record( quota_log.STAGE_RETRIEVE, unit="quota", detail={"queries": len(queries), "candidates": len(pool)} ) @@ -177,7 +275,7 @@ def cmd_search(args: argparse.Namespace) -> dict[str, Any]: raise UpstreamError( "every retrieval call failed, so the search returned nothing: " + "; ".join(dict.fromkeys(upstream_failures)), - fix=_tier1_fix(trace["stages"]["1_sources"]), + fix=_tier1_fix(trace["stages"]["1_sources"], upstream_failures), ) return { "question": args.question, @@ -260,31 +358,105 @@ def apply_rerank(pool: list[dict[str, Any]], scored: list[dict[str, Any]]) -> li ] -def _tier1_fix(sources: list[str]) -> str: - """What to actually do when discovery is down, per source. +def _tier1_fix(sources: list[str], failures: list[str] | None = None) -> str: + """What to actually do when discovery is down, per source *and per failure*. This used to say "store a Semantic Scholar API key -- it is free", which stopped being true: Ai2 no longer accept key requests from free-domain email addresses, so for a personal account that instruction has no ending. Advice that cannot be followed is worse than no advice, because it is followed first and the real fix is found second. + + The Asta branch had the same problem one step down. It said "discovery is + rate limited, not broken" whatever had happened, and pointed at a key -- but + a live run failed with `ConnectionRefusedError` raised inside Asta's own + backend, which no key affects. So the advice is now chosen by what the + failures actually say rather than assumed. """ - if sources == ["s2"]: + joined = " ".join(failures or []) + local = "Meanwhile --local-only searches what is already ingested." + elsewhere = ( + "python -m tools.paper_search search '' --tier1 {} --json " + '(or set [retrieval] tier1 = "{}" in config/grad.toml)' + ) + # The cause first where there is one, because a rate limit has a fix of its + # own and it is not the same fix as an endpoint being down. + if "rate-limited" in joined or "429" in joined: + if "asta" in sources: + return ( + "retry -- discovery is rate limited, not broken. A key raises Asta's limits " + "and is requested from a form rather than reviewed: " + f"python -m tools.jobs credential set {credentials.ASTA_KEY}. {local}" + ) + if "pwc" in sources: + return ( + "retry -- discovery is rate limited, not broken. This catalogue is " + "anonymous, so there is no key that raises it; the same literature is " + "reachable more slowly through " + elsewhere.format("asta", "asta") + f". {local}" + ) + return ( + "retry -- Semantic Scholar's anonymous pool is shared and its own keys are only " + "issued to institutional addresses. Papers with Code is anonymous and answers in " + "about a second: " + elsewhere.format("pwc", "pwc") + f". {local}" + ) + if sources == ["pwc"]: + return ( + "this catalogue is anonymous, so there is no key to add and nothing to " + "configure -- it is down or unreachable. Retry, or reach the same literature " + "more slowly through " + elsewhere.format("asta", "asta") + f". {local}" + ) + if sources in (["s2"], ["asta"]): + # Both doors onto the Semantic Scholar corpus point at the one that is + # neither rate limited nor minutes slow. return ( - "Semantic Scholar's own API only issues keys to institutional addresses, so " - "this is the shared anonymous pool. Switch to Ai2's Asta, which serves the " - "same corpus and does not require one: " - "python -m tools.paper_search search '' --tier1 asta --json " - "(or set [retrieval] tier1 = \"asta\" in config/grad.toml)" + "Semantic Scholar's own API only issues keys to institutional addresses, and " + "Asta answers in minutes rather than seconds. Papers with Code is anonymous and " + "answers in about a second: " + elsewhere.format("pwc", "pwc") + f". {local}" ) return ( - "retry -- discovery is rate limited, not broken. A key raises Asta's limits and " - "is requested from a form rather than reviewed: " - f"python -m tools.jobs credential set {credentials.ASTA_KEY}. " - "Meanwhile --local-only searches what is already ingested." + "the message above is the service's own and names the endpoint that refused -- no " + f"key affects it. The endpoints fail independently, so retrying is worth it even " + f"when one of them is down. {local}" ) +def _fill_abstracts(pool: list[dict[str, Any]], cfg: Any, trace: dict[str, Any]) -> None: + """Give the candidates that have no text an abstract, in one request. + + Stage 2 reranks on `title + snippet-or-abstract` and stage 3 triages on the + same, so a pool of bare titles is a measurably worse funnel -- and the fast + corpus does not return abstracts with its search results. Nearly every row + it does return is an arXiv paper, and arXiv takes a hundred ids at once, so + the whole pool costs one call rather than one per candidate. + + Degrades rather than fails: a candidate with no abstract ranks on its title, + which is what would have happened without this. What it must not do is go + unrecorded -- a funnel silently reranking titles looks exactly like one + reranking abstracts, and only one of them is the retrieval §5 evaluates. + """ + wanted = { + c["external"]["ArXiv"]: c + for c in pool[: http.ARXIV_BATCH] + if not (c.get("snippet") or c.get("abstract")) + and isinstance(c.get("external"), dict) + and c["external"].get("ArXiv") + } + if not wanted: + return + _progress(f"stage 1: fetching {len(wanted)} abstracts from arXiv") + found = http.arxiv_abstracts(list(wanted), cfg=cfg) + for arxiv_id, abstract in found.items(): + if arxiv_id in wanted: + wanted[arxiv_id]["abstract"] = abstract + missing = len(wanted) - len(found) + trace["stages"]["1_abstracts"] = {"wanted": len(wanted), "found": len(found)} + if missing: + trace.setdefault("warnings", []).append( + f"{missing} of {len(wanted)} candidates were reranked and triaged on their " + "title alone — arXiv returned no abstract for them" + ) + + def _local_ranked(question: str, hyde: str | None, cfg: Any, trace: dict[str, Any]) -> list[dict[str, Any]]: """Tier 2, fused across FTS5 and vectors before joining the global pool.""" path = paths.corpus_sqlite() diff --git a/ui/app.py b/ui/app.py index cf4cab1..8a140f2 100644 --- a/ui/app.py +++ b/ui/app.py @@ -47,6 +47,16 @@ ROLES = ("user", "assistant") STATIC_URL = "/grad-static" +#: How long the SDK's own interrupt is given to end the turn before the client +#: is taken down instead. The same shape as `ui/tasks.py:cancel` -- ask the thing +#: that knows how to stop cleanly, then stop it anyway -- and for the same +#: reason: a control that reports "interrupting…" and leaves the session busy +#: forever is worse than no control, because the composer then silently refuses +#: every prompt after it. +INTERRUPT_GRACE_S = 8.0 +#: How often `_stop_turn` checks whether the turn it asked to stop has settled. +SETTLE_POLL_S = 0.1 + # Where anything the transcript must not carry goes instead: this handler is the # app's own log, not user-visible text and not the persisted session file. log = logging.getLogger("grad.ui") @@ -73,6 +83,11 @@ def __init__(self, key: str = "default") -> None: # rather than at claim time means `most_recent` can tell "another window # is in this session" from "the window that was in it is gone". sessions.register(key) + #: Where a message that has nowhere else to go is put on screen. Set by + #: `build` to the workspace's status bar; None in tests and on the CLI. + #: An interrupt that failed used to be swallowed entirely, which is how + #: pressing STOP twice became the way to stop a turn. + self.notify: Any = None self.client: Any = None #: The turn in flight, as `agent.TurnStream` blocks: prose, and the tool #: calls between it. The chat window's timer draws from here, so a card @@ -89,7 +104,15 @@ def __init__(self, key: str = "default") -> None: #: Ours names the file; this one is what makes reopening continue rather #: than merely redisplay. See the note in `ui/sessions.py`. self.sdk_session_id: str | None = None - self._task: asyncio.Task[None] | None = None + #: Set while no turn is in flight, so an interrupt can wait for the turn + #: it asked to stop without polling `busy` -- which the *next* turn sets + #: back to True, and a waiter watching that flag would never wake. + self._idle = asyncio.Event() + self._idle.set() + #: The interrupt in progress, if one is. Held so the next turn can wait + #: for it: a fire-and-forget interrupt can outlive the turn it was aimed + #: at and land on the one after it. + self._stopping: asyncio.Task[None] | None = None async def start(self) -> None: if self.client is not None: @@ -234,6 +257,7 @@ async def rebind(self) -> None: await self.close() self.blocks = [] self.busy = False + self._idle.set() self.settled.clear() # Sessions live under the root's data directory, so the one that was # open belongs to the folder just left -- and the claim on it does too. @@ -252,10 +276,18 @@ async def ask(self, prompt: str, on_settle: Any) -> None: if self.busy: return self.busy = True + self._idle.clear() try: + # Before `start`, and before anything is sent. An interrupt is a + # control request to the CLI, and one still in flight is aimed at a + # turn that has already ended -- so without this it lands on the + # turn about to be issued and kills it the moment it starts, which + # is a turn that produces nothing and explains nothing. + await self._stopped() await self.start() except Exception: self.busy = False + self._idle.set() raise self.settled.append({"role": "user", "text": prompt}) @@ -271,7 +303,15 @@ async def ask(self, prompt: str, on_settle: Any) -> None: # rather than inline is what stops the two surfaces disagreeing about # whether the budget applies. result = await agent.drive_turn( - self.client, prompt, stream, session=self.session_id + self.client, + prompt, + stream, + # Recorded as it arrives rather than read off the return value, + # which an interrupted turn never reaches. That id is what lets + # the rebuilt client `resume` this conversation, so losing it on + # exactly the turns that end in a rebuild cost the whole thread. + on_session_id=self._remember_sdk_session, + session=self.session_id, ) if result.get("sdk_session_id"): self.sdk_session_id = result["sdk_session_id"] @@ -298,6 +338,10 @@ async def ask(self, prompt: str, on_settle: Any) -> None: ) finally: self.busy = False + # Woken here rather than at the end of the block: an interrupt that + # is waiting for this turn is waiting to stop *doing* things, and + # persisting and settling the transcript below are not that. + self._idle.set() # The first thing asked is what the session is about, and naming it # from the prompt beats leaving every session called "empty session" # until someone renames it by hand. An explicit title is never @@ -318,23 +362,107 @@ async def ask(self, prompt: str, on_settle: Any) -> None: self._persist() await on_settle(record) - def interrupt(self) -> None: - """Interrupt the turn in flight, if there is one. + def interrupt(self) -> str: + """Stop the turn in flight, and say what was done about it. Bound to a button and to Escape, so it fires when nothing is running. - The result has to be consumed: a bare `create_task` drops any SDK error - and Python logs "Task exception was never retrieved". + + This used to be one `create_task` around `client.interrupt()` with every + exception swallowed, and it had three failure modes that all looked the + same from the composer -- the turn stays busy, so the *next* prompt is + silently refused and nothing on screen says why. Pressing STOP a second + time appeared to fix it, which is the bug as reported: + + * the SDK refused the interrupt (not connected, control request + errored), and nothing said so; + * the interrupt was accepted but the turn did not end, and the session + stayed busy for the life of the app; + * the interrupt arrived late, after the turn had already ended, and + landed on whatever was issued next. + + So: the failure is reported, the turn is *made* to end, and the pending + interrupt is something the next turn waits for. The message is returned + rather than pushed, because the two callers already have somewhere to + put a line and disagree about where. + """ + if not self.busy: + return "nothing is running" + if self._stopping is not None and not self._stopping.done(): + return "already stopping — the turn is being taken down" + self._stopping = asyncio.create_task(self._stop_turn()) + return "interrupting the turn…" + + async def _stop_turn(self) -> None: + """Ask the SDK to stop, then make sure the turn actually stopped. + + The escalation is `ui/tasks.py:cancel`'s, for the same reason: the tool's + own stop verb is the one that ends things cleanly, and the blunt + instrument is what happens when it does not work. Here the blunt + instrument is dropping the client, which ends `receive_response` and lets + `ask` settle the partial turn the way any other failure settles. + + **The client is dropped either way**, and that is deliberate rather than + laziness about the happy path. The client owns the CLI subprocess and the + message stream, and an interrupted turn is precisely the case where what + is left in that stream is unclear -- a result the aborted turn never + emitted, or one nobody read. Reusing it made the *next* turn end + instantly on a message belonging to the last one, with nothing drawn. + A fresh client cannot have a stale message in it, and `resume` carries + the conversation across, so what is paid is a restart on a deliberate, + occasional action. """ - if not (self.busy and self.client and hasattr(self.client, "interrupt")): + client = self.client + if client is not None and hasattr(client, "interrupt"): + try: + await client.interrupt() + except Exception as exc: # noqa: BLE001 - reported, never swallowed + log.warning("the SDK refused the interrupt", exc_info=exc) + self.say( + f"the SDK refused the interrupt ({type(exc).__name__}) — " + "stopping the turn the hard way" + ) + if not await self._idles_within(INTERRUPT_GRACE_S): + self.say(f"the turn did not stop within {INTERRUPT_GRACE_S:.0f}s — taking the client down") + # Whether it stopped when asked or had to be taken down, the client goes. + await self.close() + if self.busy and not await self._idles_within(INTERRUPT_GRACE_S): + # The turn outlived the client that was feeding it. Whatever it is + # waiting for is not going to arrive, and a composer that stays + # locked on the outcome of that is the failure this whole method is + # about -- so the flag is cleared and the fact is said out loud + # rather than left to be inferred from a session that never answers. + self.busy = False + self._idle.set() + self.say("the turn is still winding down — the composer is usable again") + + async def _idles_within(self, seconds: float) -> bool: + """Has the turn settled inside this window? Never raises.""" + try: + await asyncio.wait_for(self._idle.wait(), timeout=seconds) + except asyncio.TimeoutError: + return False + return True + + async def _stopped(self) -> None: + """Wait for a pending interrupt, so it cannot land on the next turn.""" + pending, self._stopping = self._stopping, None + if pending is None or pending.done(): return + # Bounded: `_stop_turn` bounds itself, and a hang here would be the very + # thing this method exists to prevent, one layer up. + await asyncio.wait([pending], timeout=INTERRUPT_GRACE_S * 2) - async def _interrupt() -> None: - try: - await self.client.interrupt() - except Exception: # noqa: BLE001 - a failed interrupt must not kill the app - pass + def _remember_sdk_session(self, sdk_session_id: str) -> None: + self.sdk_session_id = sdk_session_id - self._task = asyncio.create_task(_interrupt()) + def say(self, message: str) -> None: + """A line for the status bar, when there is one to put it in.""" + log.info("%s", message) + if self.notify is not None: + try: + self.notify(message) + except Exception: # noqa: BLE001 - a notice must not kill a turn + log.exception("could not report: %s", message) def path(self) -> Path: return sessions.path_for(self.session_id) @@ -442,6 +570,11 @@ def index() -> None: session = Session(_client_key()) session.adopt() workspace = state_mod.Workspace(session, _current_project()) + # The one place both exist. A turn's own failures already reach the + # transcript; this is for the things that happen *around* a turn -- an + # interrupt the SDK refused, a client that had to be taken down -- which + # have no turn to be written into. + session.notify = workspace.say # Per client, not `app.on_shutdown`: that would accumulate one handler # per connection and hold every session's subprocess open until the app # itself exits. diff --git a/ui/kit.py b/ui/kit.py index 4369423..d07c649 100644 --- a/ui/kit.py +++ b/ui/kit.py @@ -313,6 +313,93 @@ def blink_caret() -> Any: return el("span", "grad-caret") +class Menu: + """A dialog whose body is rebuilt each time it opens. + + `ui.dialog` builds its contents once. These menus list projects, folders, + open windows and stored sessions, and all four change *because of* what the + dialog does -- create a project and the list it was read from is already + stale, open a window and the mark beside its name is wrong. Redrawing on open + is cheaper than binding every row to the poll, and it cannot go stale between + the click and the dialog appearing. + + `draw` is handed the menu so a control *inside* it can call `redraw` after + changing what the menu is listing -- which is what lets the window menu stay + open across several toggles instead of closing after each one. + + It lives here rather than in `ui/shell.py`, where it was written, because the + chat window's session picker is the fourth of these: a Quasar `select` was + the one control in the workspace still carrying NiceGUI's own look, and it + could not say the two things a session row has to say -- that reopening one + only redisplays it, and that another window already has it. + """ + + def __init__(self, dialog: Any, draw: Any) -> None: + self._dialog = dialog + self._draw = draw + + def open(self) -> None: + self.redraw() + self._dialog.open() + + def redraw(self) -> None: + self._draw(self) + + def close(self) -> None: + self._dialog.close() + + +def menu(draw: Callable[[Any, Any], None], *, width: int = 460) -> Menu: + """A `Menu` over a card in the app's own paper, ready to be filled. + + `draw` is called with `(body, menu)` each time it opens; it is expected to + clear the body itself, because a redraw from inside the menu is the same + call. + """ + ui = _ui() + with ui.dialog() as dialog, el("div", "grad-app"): + body = el("div", "grad-card", style=f"background: var(--grad-paper); min-width: {width}px") + return Menu(dialog, lambda m: draw(body, m)) + + +def menu_row( + mark: str, + name: str, + hint: str, + *, + open: bool = False, + title: str = "", + wide: bool = False, + disabled: bool = False, +) -> Any: + """One row of a menu: a mark, a name, and what the row is for. + + A `