drydock-cli 3.0.46__tar.gz → 3.0.49__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (65) hide show
  1. {drydock_cli-3.0.46/drydock_cli.egg-info → drydock_cli-3.0.49}/PKG-INFO +1 -1
  2. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/drydock/agent.py +4 -2
  3. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/drydock/cli.py +46 -5
  4. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/drydock/config.py +6 -0
  5. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/drydock/detect.py +1 -1
  6. drydock_cli-3.0.49/drydock/graphrag.py +244 -0
  7. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/drydock/loop_detect.py +43 -0
  8. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/drydock/providers.py +34 -4
  9. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/drydock/tools/__init__.py +56 -3
  10. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/drydock/tui/app.py +83 -6
  11. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/drydock/tui/messages.py +6 -0
  12. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/drydock/tui/widgets.py +24 -1
  13. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/drydock/tuning.py +27 -0
  14. {drydock_cli-3.0.46 → drydock_cli-3.0.49/drydock_cli.egg-info}/PKG-INFO +1 -1
  15. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/drydock_cli.egg-info/SOURCES.txt +4 -0
  16. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/pyproject.toml +1 -1
  17. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_config.py +4 -1
  18. drydock_cli-3.0.49/tests/test_context_limit_config.py +44 -0
  19. drydock_cli-3.0.49/tests/test_graphrag.py +75 -0
  20. drydock_cli-3.0.49/tests/test_runaway_repetition.py +128 -0
  21. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_tuning.py +54 -0
  22. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/LICENSE +0 -0
  23. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/NOTICE +0 -0
  24. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/README.md +0 -0
  25. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/drydock/__init__.py +0 -0
  26. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/drydock/__main__.py +0 -0
  27. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/drydock/bash_safety.py +0 -0
  28. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/drydock/compaction.py +0 -0
  29. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/drydock/guards.py +0 -0
  30. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/drydock/tool_registry.py +0 -0
  31. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/drydock/tui/__init__.py +0 -0
  32. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/drydock/tui/approval.py +0 -0
  33. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/drydock_cli.egg-info/dependency_links.txt +0 -0
  34. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/drydock_cli.egg-info/entry_points.txt +0 -0
  35. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/drydock_cli.egg-info/requires.txt +0 -0
  36. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/drydock_cli.egg-info/top_level.txt +0 -0
  37. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/setup.cfg +0 -0
  38. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_approval.py +0 -0
  39. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_back_command.py +0 -0
  40. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_bash_process_group.py +0 -0
  41. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_bash_safety.py +0 -0
  42. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_bash_timeout_network.py +0 -0
  43. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_cli_agents.py +0 -0
  44. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_compact_command.py +0 -0
  45. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_compaction.py +0 -0
  46. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_config_migration.py +0 -0
  47. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_detect.py +0 -0
  48. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_empty_response.py +0 -0
  49. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_failure_loop.py +0 -0
  50. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_first_run_setup.py +0 -0
  51. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_guards_and_tools.py +0 -0
  52. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_hallucinated_tools.py +0 -0
  53. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_leaked_tool_call.py +0 -0
  54. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_loop_detect.py +0 -0
  55. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_oneshot_unreachable.py +0 -0
  56. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_plan_autocontinue.py +0 -0
  57. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_providers_unreachable.py +0 -0
  58. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_stop.py +0 -0
  59. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_streaming_newlines.py +0 -0
  60. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_subagent.py +0 -0
  61. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_todo.py +0 -0
  62. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_tool_arg_parsing.py +0 -0
  63. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_tools_undo.py +0 -0
  64. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_tui.py +0 -0
  65. {drydock_cli-3.0.46 → drydock_cli-3.0.49}/tests/test_vision_input.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: drydock-cli
3
- Version: 3.0.46
3
+ Version: 3.0.49
4
4
  Summary: Drydock — a local, provider-agnostic terminal coding agent for local LLMs
5
5
  Author: Frank Bobe III
6
6
  License-Expression: Apache-2.0
@@ -19,7 +19,7 @@ def _plan_has_unfinished(config: dict) -> bool:
19
19
  todo = config.get("_todo")
20
20
  return bool(todo) and any(status != "done" for _, status in todo)
21
21
 
22
- from drydock.providers import stream, AssistantTurn, TextChunk
22
+ from drydock.providers import stream, AssistantTurn, ReasoningChunk, TextChunk
23
23
  from drydock.tool_registry import schemas, execute
24
24
  from drydock.tools import register_all
25
25
 
@@ -168,7 +168,9 @@ def run(
168
168
  tool_schemas=filter_tool_schemas(available, turn_config.get("model")),
169
169
  config=turn_config,
170
170
  ):
171
- if isinstance(event, TextChunk):
171
+ if isinstance(event, ReasoningChunk):
172
+ yield event
173
+ elif isinstance(event, TextChunk):
172
174
  yield event
173
175
  elif isinstance(event, AssistantTurn):
174
176
  assistant_turn = event
@@ -149,7 +149,8 @@ def run_oneshot(prompt: str, config: dict) -> None:
149
149
 
150
150
  def handle_command(cmd: str, state: AgentState, config: dict) -> bool:
151
151
  """Handle slash commands. Returns True if handled."""
152
- cmd = cmd.lower().strip()
152
+ raw = cmd.strip() # case-preserved (paths/args must not be lowercased)
153
+ cmd = raw.lower()
153
154
  if cmd == "/help":
154
155
  print_colored("Commands:", "bold")
155
156
  print(" /help — show this help")
@@ -164,7 +165,7 @@ def handle_command(cmd: str, state: AgentState, config: dict) -> bool:
164
165
  print_colored(" Conversation cleared.", "green")
165
166
  return True
166
167
  elif cmd == "/compact":
167
- from drydock.compaction import compact, estimate_tokens
168
+ from drydock.compaction import compact, emergency_compact, estimate_tokens
168
169
 
169
170
  if not state.messages:
170
171
  print_colored(" Nothing to compact.", "dim")
@@ -172,6 +173,8 @@ def handle_command(cmd: str, state: AgentState, config: dict) -> bool:
172
173
  before = estimate_tokens(state.messages)
173
174
  limit = config.get("context_limit", 65536) or 65536
174
175
  state.messages = compact(state.messages, limit)
176
+ if estimate_tokens(state.messages) > limit * 0.5:
177
+ state.messages = emergency_compact(state.messages, limit)
175
178
  after = estimate_tokens(state.messages)
176
179
  saved = before - after
177
180
  if saved > 0:
@@ -187,6 +190,37 @@ def handle_command(cmd: str, state: AgentState, config: dict) -> bool:
187
190
  print_colored(f" Messages: {len(state.messages)}", "cyan")
188
191
  print_colored(f" Tokens: {state.total_input_tokens}in / {state.total_output_tokens}out", "cyan")
189
192
  return True
193
+ elif cmd.split()[0] == "/graphrag":
194
+ from drydock import graphrag
195
+
196
+ cwd = config.get("cwd") or os.getcwd()
197
+ store = graphrag.default_store_path(cwd)
198
+ parts = raw.split() # raw = case-preserved, so paths survive
199
+ sub = parts[1].lower() if len(parts) > 1 else ""
200
+ rest = " ".join(parts[2:]).strip()
201
+ if sub == "build" and rest:
202
+ print_colored(f" Building knowledge base from {rest} …", "dim")
203
+ try:
204
+ stats = graphrag.build_index([rest], store, cwd=cwd)
205
+ print_colored(
206
+ f" ✓ {stats['files']} files · {stats['chunks']} chunks · "
207
+ f"{stats['entities']} entities · {stats['edges']} edges → {store}",
208
+ "green",
209
+ )
210
+ except Exception as e: # noqa: BLE001
211
+ print_colored(f" graphrag build failed: {e}", "red")
212
+ elif sub == "clear":
213
+ store.unlink(missing_ok=True)
214
+ print_colored(" Knowledge base cleared.", "green")
215
+ else:
216
+ index = graphrag.load_index(store)
217
+ if index is None:
218
+ print_colored(" No knowledge base. Build: /graphrag build <path>", "dim")
219
+ else:
220
+ print_colored(
221
+ f" Knowledge base: {len(index.get('chunks', []))} chunks · "
222
+ f"{len(index.get('entities', {}))} entities ({store})", "cyan")
223
+ return True
190
224
  elif cmd in ("/quit", "/exit"):
191
225
  raise KeyboardInterrupt
192
226
  return False
@@ -246,6 +280,11 @@ def main():
246
280
  parser.add_argument("--provider", help="Provider: vllm, ollama, lmstudio, openai")
247
281
  parser.add_argument("--base-url", dest="base_url", help="Override API base URL")
248
282
  parser.add_argument("--max-tokens", dest="max_tokens", type=int, help="Max response tokens")
283
+ parser.add_argument(
284
+ "--context-limit", dest="context_limit", type=int,
285
+ help="Model server context window (must match the server's -c / "
286
+ "--max-model-len; drives the ctx gauge + when compaction fires)",
287
+ )
249
288
  parser.add_argument("--temperature", type=float)
250
289
  parser.add_argument("--max-tool-calls", type=int, default=0, help="Max tool calls (0=unlimited)")
251
290
  parser.add_argument("--force-first-tool", action="store_true", help="Force tool_choice=required on first turn")
@@ -268,6 +307,7 @@ def main():
268
307
  "base_url": args.base_url,
269
308
  "max_tokens": args.max_tokens,
270
309
  "temperature": args.temperature,
310
+ "context_limit": args.context_limit,
271
311
  }, cfg_path)
272
312
 
273
313
  # First launch: ask the user for their model server URL + model name and
@@ -288,10 +328,11 @@ def main():
288
328
  cfgmod.save_file(cfg, cfg_path)
289
329
 
290
330
  config = {
331
+ # context_limit now comes from cfg (DEFAULTS < config.toml < --context-limit)
332
+ # — it drives the ctx gauge AND when compaction fires, so it must match the
333
+ # server's real -c. Previously hardcoded to 65536 here, which clobbered the
334
+ # config value and overflowed servers running a smaller window.
291
335
  **cfg,
292
- # Match the llama.cpp server's -c (gemma4 runs at 64k). Drives both the
293
- # status-bar context gauge and when compaction kicks in.
294
- "context_limit": 65536,
295
336
  "max_tool_calls": args.max_tool_calls,
296
337
  "force_first_tool": args.force_first_tool,
297
338
  "_approve_all": args.dangerously_skip_permissions,
@@ -28,6 +28,12 @@ DEFAULTS: dict[str, object] = {
28
28
  "base_url": "http://localhost:8000/v1",
29
29
  "max_tokens": 8192, # 4096 truncated large file writes mid-JSON (→ _raw fail)
30
30
  "temperature": 0.2,
31
+ # The model server's context window (llama.cpp -c / vLLM --max-model-len).
32
+ # Drives the ctx gauge AND when compaction fires (it compacts at ~60% of
33
+ # this). It MUST match your server: if it's too high, drydock overflows the
34
+ # real window BEFORE compacting (the server 400s) — set this to your -c
35
+ # value. Default 65536 matches the bundled gemma4 server (-c 65536).
36
+ "context_limit": 65536,
31
37
  "theme": "harbor",
32
38
  }
33
39
 
@@ -63,7 +63,7 @@ def onboarding_message(found: list[dict]) -> str:
63
63
  return (
64
64
  "No local LLM detected (" + ports + "). Point Drydock at your model "
65
65
  "right here:\n"
66
- " /model url <http://host:port/v1> then /model <model-name>\n"
66
+ " /model url <http://localhost:8000/v1> then /model <model-name>\n"
67
67
  "(saved to ~/.drydock/config.toml). Or start a local server "
68
68
  "(llama.cpp / vLLM / Ollama / LM Studio) and restart."
69
69
  )
@@ -0,0 +1,244 @@
1
+ """GraphRAG — a local, dependency-free knowledge graph for retrieval.
2
+
3
+ The user builds an index from their own docs/code; the agent then queries it
4
+ through the `Knowledge` tool while answering or coding, so it can ground its
5
+ work in project-specific information the model was never trained on.
6
+
7
+ Design (clean-room, stdlib only — no embeddings, no third-party deps, works
8
+ against any server):
9
+ • Ingest text files → split into chunks.
10
+ • Extract ENTITIES from each chunk (proper nouns, `code` terms, CamelCase /
11
+ snake_case identifiers, ACRONYMS) → graph NODES.
12
+ • Entities co-occurring in a chunk get an EDGE (weighted by co-occurrence) →
13
+ the graph. Each entity also points at the chunks it appears in.
14
+ • Query: pull the query's entities + keywords, score chunks by entity and
15
+ keyword overlap, then EXPAND one hop along the entity graph so strongly
16
+ related context comes along — the "graph" in GraphRAG. Return the top
17
+ chunks plus the related entities, both fed back to the model.
18
+
19
+ The index is a single JSON file so it's inspectable, portable, and gitignorable.
20
+
21
+ All logic original to Drydock.
22
+ """
23
+ from __future__ import annotations
24
+
25
+ import json
26
+ import os
27
+ import re
28
+ from collections import Counter, defaultdict
29
+ from pathlib import Path
30
+
31
+ # Files we ingest as text. Everything else (binaries, images) is skipped.
32
+ _TEXT_EXT = {
33
+ ".md", ".txt", ".rst", ".py", ".js", ".ts", ".tsx", ".jsx", ".java", ".go",
34
+ ".rs", ".c", ".h", ".cpp", ".hpp", ".cs", ".rb", ".php", ".sh", ".toml",
35
+ ".yaml", ".yml", ".json", ".cfg", ".ini", ".sql", ".html", ".css", ".org",
36
+ ".tex", ".csv", ".log", ".",
37
+ }
38
+ _SKIP_DIRS = {".git", "node_modules", "__pycache__", ".venv", "venv", "dist",
39
+ "build", ".mypy_cache", ".pytest_cache", ".drydock"}
40
+
41
+ _STOPWORDS = frozenset("""
42
+ the a an and or but if then else for while of to in on at by with from into as
43
+ is are was were be been being do does did this that these those it its their
44
+ your you we they he she his her our not no yes can will would should could may
45
+ """.split())
46
+
47
+ # Entity patterns (kept simple + fast; tuned to surface useful nodes, not noise)
48
+ _RE_BACKTICK = re.compile(r"`([^`\n]{2,60})`")
49
+ _RE_IDENT = re.compile(r"\b([a-zA-Z_][a-zA-Z0-9_]*(?:[_./][a-zA-Z0-9_]+)+)\b")
50
+ _RE_CAMEL = re.compile(r"\b([A-Z][a-z0-9]+(?:[A-Z][a-z0-9]+)+)\b")
51
+ _RE_PROPER = re.compile(r"\b([A-Z][a-zA-Z0-9]+(?:[ \t]+[A-Z][a-zA-Z0-9]+){0,3})\b")
52
+ _RE_ACRONYM = re.compile(r"\b([A-Z]{2,6})\b")
53
+ _RE_WORD = re.compile(r"[a-zA-Z0-9_]+")
54
+
55
+
56
+ def default_store_path(cwd: str) -> Path:
57
+ """Project-local index (travels with the project, easy to .gitignore)."""
58
+ return Path(cwd) / ".drydock" / "graphrag.json"
59
+
60
+
61
+ def extract_entities(text: str) -> list[str]:
62
+ """Heuristic entity extraction → normalized (lowercased) entity keys."""
63
+ found: set[str] = set()
64
+ for rx in (_RE_BACKTICK, _RE_IDENT, _RE_CAMEL, _RE_PROPER, _RE_ACRONYM):
65
+ for m in rx.findall(text):
66
+ e = m.strip().lower()
67
+ # Drop trivial / stopword-only entities and overlong noise.
68
+ if len(e) < 3 or len(e) > 60:
69
+ continue
70
+ if all(w in _STOPWORDS for w in _RE_WORD.findall(e)):
71
+ continue
72
+ found.add(e)
73
+ return sorted(found)
74
+
75
+
76
+ def _chunk_text(text: str, max_chars: int = 1200) -> list[str]:
77
+ """Split into paragraph-ish chunks, capping size so a chunk stays focused."""
78
+ chunks: list[str] = []
79
+ buf: list[str] = []
80
+ size = 0
81
+ for para in re.split(r"\n\s*\n", text):
82
+ para = para.strip()
83
+ if not para:
84
+ continue
85
+ if size + len(para) > max_chars and buf:
86
+ chunks.append("\n\n".join(buf))
87
+ buf, size = [], 0
88
+ if len(para) > max_chars: # a single huge paragraph → hard-split
89
+ for i in range(0, len(para), max_chars):
90
+ chunks.append(para[i:i + max_chars])
91
+ continue
92
+ buf.append(para)
93
+ size += len(para)
94
+ if buf:
95
+ chunks.append("\n\n".join(buf))
96
+ return chunks
97
+
98
+
99
+ def _iter_text_files(paths: list[str]):
100
+ for p in paths:
101
+ path = Path(p)
102
+ if path.is_file():
103
+ yield path
104
+ elif path.is_dir():
105
+ for root, dirs, files in os.walk(path):
106
+ dirs[:] = [d for d in dirs if d not in _SKIP_DIRS]
107
+ for f in files:
108
+ fp = Path(root) / f
109
+ if fp.suffix.lower() in _TEXT_EXT or fp.suffix == "":
110
+ yield fp
111
+
112
+
113
+ def build_index(paths: list[str], store_path: str | Path, *, cwd: str = ".") -> dict:
114
+ """Build (or rebuild) the knowledge graph from paths; persist to store_path.
115
+
116
+ Returns a stats dict: {files, chunks, entities, edges}.
117
+ """
118
+ chunks: list[dict] = []
119
+ entity_chunks: dict[str, set[int]] = defaultdict(set)
120
+ edges: dict[str, Counter] = defaultdict(Counter)
121
+ files = 0
122
+
123
+ for fp in _iter_text_files([str(Path(cwd) / p) if not os.path.isabs(p) else p
124
+ for p in paths]):
125
+ try:
126
+ text = fp.read_text("utf-8", "ignore")
127
+ except OSError:
128
+ continue
129
+ if not text.strip():
130
+ continue
131
+ files += 1
132
+ rel = os.path.relpath(str(fp), cwd)
133
+ for body in _chunk_text(text):
134
+ cid = len(chunks)
135
+ ents = extract_entities(body)
136
+ chunks.append({"id": cid, "source": rel, "text": body, "entities": ents})
137
+ for e in ents:
138
+ entity_chunks[e].add(cid)
139
+ for i, a in enumerate(ents): # co-occurrence edges within the chunk
140
+ for b in ents[i + 1:]:
141
+ edges[a][b] += 1
142
+ edges[b][a] += 1
143
+
144
+ index = {
145
+ "version": 1,
146
+ "chunks": chunks,
147
+ "entities": {e: sorted(cids) for e, cids in entity_chunks.items()},
148
+ "edges": {a: dict(nbrs) for a, nbrs in edges.items()},
149
+ }
150
+ sp = Path(store_path)
151
+ sp.parent.mkdir(parents=True, exist_ok=True)
152
+ sp.write_text(json.dumps(index), encoding="utf-8")
153
+ return {
154
+ "files": files,
155
+ "chunks": len(chunks),
156
+ "entities": len(entity_chunks),
157
+ "edges": sum(len(v) for v in edges.values()) // 2,
158
+ }
159
+
160
+
161
+ def load_index(store_path: str | Path) -> dict | None:
162
+ sp = Path(store_path)
163
+ if not sp.exists():
164
+ return None
165
+ try:
166
+ return json.loads(sp.read_text("utf-8"))
167
+ except (OSError, ValueError):
168
+ return None
169
+
170
+
171
+ def query_index(index: dict, query: str, *, k: int = 5, hops: bool = True) -> dict:
172
+ """Retrieve the top-k chunks for a query, expanded one hop over the graph.
173
+
174
+ Returns {chunks: [{source, text, score}], related: [entities]}.
175
+ """
176
+ chunks = index.get("chunks", [])
177
+ entities = index.get("entities", {})
178
+ edges = index.get("edges", {})
179
+ if not chunks:
180
+ return {"chunks": [], "related": []}
181
+
182
+ q_entities = set(extract_entities(query))
183
+ q_words = {w for w in (w.lower() for w in _RE_WORD.findall(query))
184
+ if len(w) > 2 and w not in _STOPWORDS}
185
+
186
+ # Match query entities to graph entities (exact + substring both ways).
187
+ matched: set[str] = set()
188
+ for qe in q_entities | q_words:
189
+ for ge in entities:
190
+ if qe == ge or qe in ge or ge in qe:
191
+ matched.add(ge)
192
+
193
+ # 1-hop graph expansion: pull the strongest neighbors of matched entities.
194
+ related: list[str] = []
195
+ expanded: set[str] = set(matched)
196
+ if hops:
197
+ for e in list(matched):
198
+ nbrs = sorted(edges.get(e, {}).items(), key=lambda kv: -kv[1])[:5]
199
+ for nbr, _w in nbrs:
200
+ expanded.add(nbr)
201
+ if nbr not in matched:
202
+ related.append(nbr)
203
+
204
+ # Score chunks: direct entity hits weigh most, then graph-neighbor hits,
205
+ # then raw keyword overlap (so retrieval still works if entities miss).
206
+ scores: dict[int, float] = defaultdict(float)
207
+ for e in matched:
208
+ for cid in entities.get(e, []):
209
+ scores[cid] += 3.0
210
+ for e in expanded - matched:
211
+ for cid in entities.get(e, []):
212
+ scores[cid] += 1.0
213
+ for c in chunks:
214
+ if q_words:
215
+ overlap = len(q_words & {w.lower() for w in _RE_WORD.findall(c["text"])})
216
+ if overlap:
217
+ scores[c["id"]] += 0.5 * overlap
218
+
219
+ ranked = sorted(scores.items(), key=lambda kv: -kv[1])[:k]
220
+ out_chunks = [
221
+ {"source": chunks[cid]["source"], "text": chunks[cid]["text"], "score": round(sc, 1)}
222
+ for cid, sc in ranked if sc > 0
223
+ ]
224
+ # de-dup related, keep order, cap
225
+ seen: set[str] = set()
226
+ rel_unique = [r for r in related if not (r in seen or seen.add(r))][:12]
227
+ return {"chunks": out_chunks, "related": rel_unique}
228
+
229
+
230
+ def format_results(result: dict, query: str) -> str:
231
+ """Render query results as the tool's string output for the model."""
232
+ chunks = result.get("chunks", [])
233
+ if not chunks:
234
+ return (
235
+ f"No knowledge-base matches for: {query!r}. The index may not cover "
236
+ "this topic — answer from your own knowledge, or ask the user to add "
237
+ "the relevant docs with /graphrag build <path>."
238
+ )
239
+ parts = [f"Knowledge base — {len(chunks)} relevant passage(s) for {query!r}:"]
240
+ for i, c in enumerate(chunks, 1):
241
+ parts.append(f"\n[{i}] (source: {c['source']}, score {c['score']})\n{c['text']}")
242
+ if result.get("related"):
243
+ parts.append("\nRelated entities in the graph: " + ", ".join(result["related"]))
244
+ return "\n".join(parts)
@@ -86,6 +86,49 @@ def path_thrash_note(path: str, count: int) -> str | None:
86
86
  )
87
87
 
88
88
 
89
+ def runaway_repetition_len(
90
+ text: str,
91
+ *,
92
+ window: int = 2000,
93
+ min_run: int = 600,
94
+ max_unit: int = 80,
95
+ min_reps: int = 6,
96
+ ) -> int:
97
+ """Length (in chars) of a runaway repeated SUFFIX of *text*, else 0.
98
+
99
+ A weak local model sometimes collapses mid-generation into emitting the same
100
+ short unit hundreds of times (observed: gemma4 streamed ``295:`` ~1365× on
101
+ make-mips-interpreter). The tool-call loop guard never sees this — it's
102
+ streamed assistant TEXT, not a tool call — so the turn balloons unchecked.
103
+
104
+ This detects that one failure mode and nothing else. It is deliberately
105
+ CONSERVATIVE: it requires a long pure-repetition run (>= ``min_run`` chars
106
+ AND >= ``min_reps`` repeats of one <= ``max_unit``-char unit) so that
107
+ legitimate repetition — a ``---`` rule, a short bulleted list, an ASCII
108
+ table — never trips it. Pure-whitespace units (the blank line between
109
+ markdown blocks) are ignored. Returns the run length so the caller can trim
110
+ exactly the repeated tail; 0 means "looks fine, keep going".
111
+ """
112
+ if len(text) < min_run:
113
+ return 0
114
+ tail = text[-window:]
115
+ n = len(tail)
116
+ best = 0
117
+ for p in range(1, min(max_unit, n // min_reps) + 1):
118
+ unit = tail[-p:]
119
+ if not unit.strip():
120
+ continue # ignore blank-line / whitespace units
121
+ reps = 0
122
+ i = n
123
+ while i - p >= 0 and tail[i - p:i] == unit:
124
+ reps += 1
125
+ i -= p
126
+ run = reps * p
127
+ if reps >= min_reps and run >= min_run:
128
+ best = max(best, run)
129
+ return best
130
+
131
+
89
132
  class LoopTracker:
90
133
  """Counts identical tool calls across a run and produces advisory notes."""
91
134
 
@@ -23,7 +23,8 @@ _LLM_POOL = concurrent.futures.ThreadPoolExecutor(max_workers=8, thread_name_pre
23
23
  class _StopRequested(Exception):
24
24
  """Internal: STOP was pressed during a blocking LLM call."""
25
25
 
26
- from drydock.tuning import strip_leaked_tool_calls, strip_thinking_tokens, use_streaming
26
+ from drydock.loop_detect import runaway_repetition_len
27
+ from drydock.tuning import extract_thinking, strip_leaked_tool_calls, strip_thinking_tokens, use_streaming
27
28
 
28
29
  # ── Provider registry ─────────────────────────────────────────────────────
29
30
 
@@ -167,6 +168,14 @@ def _stream_chunks(response, cancel):
167
168
  class TextChunk:
168
169
  text: str
169
170
 
171
+ @dataclass
172
+ class ReasoningChunk:
173
+ """Thinking tokens emitted by the model before its answer (Gemma 4 only).
174
+ Yielded once per non-streaming turn, before any TextChunk, when the model
175
+ included a <|channel>…<channel|> block in its response.
176
+ """
177
+ text: str
178
+
170
179
  @dataclass
171
180
  class AssistantTurn:
172
181
  text: str
@@ -405,6 +414,11 @@ def stream(
405
414
  text = ""
406
415
  tool_buf: dict = {} # index → {id, name, args}
407
416
  in_tok = out_tok = 0
417
+ # Runaway-repetition guard: a weak model can collapse into streaming one
418
+ # short unit hundreds of times (gemma4: `295:` ×1365). Advisory, never
419
+ # raises — on detection we trim the repeated tail, emit a note, and stop
420
+ # reading (the orphaned producer drains in the background, like STOP).
421
+ _rep_checked_at = 0
408
422
 
409
423
  try:
410
424
  response = _create_abortable(client, kwargs, base_url, provider, cancel, read_timeout)
@@ -432,6 +446,20 @@ def stream(
432
446
  if chunk_text:
433
447
  text += chunk_text
434
448
  yield TextChunk(chunk_text)
449
+ # Throttled: only scan once text is long enough and every ~200
450
+ # new chars, so the common path pays almost nothing.
451
+ if len(text) - _rep_checked_at >= 200:
452
+ _rep_checked_at = len(text)
453
+ run = runaway_repetition_len(text)
454
+ if run:
455
+ text = (
456
+ text[: len(text) - run]
457
+ + "\n[… output began repeating — stopped by drydock]"
458
+ )
459
+ yield TextChunk(
460
+ "\n[stopped — model output began repeating itself]"
461
+ )
462
+ break
435
463
 
436
464
  if delta.tool_calls:
437
465
  for tc in delta.tool_calls:
@@ -479,10 +507,12 @@ def _complete_nonstreaming(
479
507
  msg = choice.message
480
508
 
481
509
  # Order matters: remove <|tool_call> blobs FIRST (sets had_leak for the
482
- # retry nudge), THEN strip channel/special-token markers — otherwise the
483
- # generic special-token pass would eat the tool_call markers.
510
+ # retry nudge), THEN extract+strip channel/special-token markers — otherwise
511
+ # the generic special-token pass would eat the tool_call markers.
484
512
  text, had_leak = strip_leaked_tool_calls(msg.content or "")
485
- text = strip_thinking_tokens(text)
513
+ thinking, text = extract_thinking(text)
514
+ if thinking:
515
+ yield ReasoningChunk(thinking)
486
516
  if text.strip():
487
517
  yield TextChunk(text)
488
518
 
@@ -161,6 +161,32 @@ SCHEMAS = [
161
161
  "required": ["prompt"],
162
162
  },
163
163
  },
164
+ {
165
+ "name": "Knowledge",
166
+ "description": (
167
+ "Search the user's KNOWLEDGE BASE (a GraphRAG index they built from "
168
+ "their own docs/code) for project-specific information you were not "
169
+ "trained on. Returns the most relevant passages plus related entities "
170
+ "from the graph. Use it BEFORE answering or coding when the task may "
171
+ "depend on the user's private/project knowledge (their APIs, specs, "
172
+ "data, conventions). If it returns no matches, the topic isn't in the "
173
+ "base — answer normally."
174
+ ),
175
+ "input_schema": {
176
+ "type": "object",
177
+ "properties": {
178
+ "query": {
179
+ "type": "string",
180
+ "description": "What to look up, as a natural-language question or keywords.",
181
+ },
182
+ "k": {
183
+ "type": "integer",
184
+ "description": "Max passages to return (default 5).",
185
+ },
186
+ },
187
+ "required": ["query"],
188
+ },
189
+ },
164
190
  ]
165
191
 
166
192
  # ── Tool implementations ──────────────────────────────────────────────────
@@ -709,6 +735,32 @@ def tool_task(params: dict, config: dict) -> str:
709
735
  return f"[sub-agent finished {steps} step(s) with no summary]"
710
736
 
711
737
 
738
+ def tool_knowledge(params: dict, config: dict) -> str:
739
+ """Query the project's GraphRAG knowledge base (built with /graphrag build).
740
+ Read-only; returns the most relevant passages plus related graph entities.
741
+ If no index exists, says so cleanly rather than erroring."""
742
+ from drydock import graphrag
743
+
744
+ query = (params.get("query") or "").strip()
745
+ if not query:
746
+ return "Error: `Knowledge` needs a `query` describing what to look up."
747
+ cwd = config.get("cwd") or os.getcwd()
748
+ store = config.get("graphrag_store") or graphrag.default_store_path(cwd)
749
+ index = graphrag.load_index(store)
750
+ if index is None:
751
+ return (
752
+ "No knowledge base has been built yet. The user can build one with "
753
+ "'/graphrag build <path>' (a file or directory of docs/code). Until "
754
+ "then, answer from your own knowledge."
755
+ )
756
+ try:
757
+ k = int(params.get("k") or 5)
758
+ except (TypeError, ValueError):
759
+ k = 5
760
+ result = graphrag.query_index(index, query, k=max(1, min(k, 15)))
761
+ return graphrag.format_results(result, query)
762
+
763
+
712
764
  # ── Register all tools ────────────────────────────────────────────────────
713
765
 
714
766
  _TOOLS = [
@@ -720,6 +772,7 @@ _TOOLS = [
720
772
  ("Grep", tool_grep, True),
721
773
  ("todo", tool_todo, False),
722
774
  ("task", tool_task, True),
775
+ ("Knowledge", tool_knowledge, True),
723
776
  ]
724
777
 
725
778
  def register_all():
@@ -728,10 +781,10 @@ def register_all():
728
781
  func = {
729
782
  "Read": tool_read, "Write": tool_write, "Edit": tool_edit,
730
783
  "Bash": tool_bash, "Glob": tool_glob, "Grep": tool_grep,
731
- "todo": tool_todo, "task": tool_task,
784
+ "todo": tool_todo, "task": tool_task, "Knowledge": tool_knowledge,
732
785
  }[name]
733
- # `task` is read-only w.r.t. the parent's files (it can't Write/Edit).
734
- read_only = name in ("Read", "Glob", "Grep", "task")
786
+ # task + Knowledge are read-only w.r.t. the parent's files.
787
+ read_only = name in ("Read", "Glob", "Grep", "task", "Knowledge")
735
788
  register(ToolDef(name=name, schema=schema, func=func, read_only=read_only))
736
789
 
737
790
  register_all()