tokensafe 1.0.1 → 1.1.0

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.
package/README.md CHANGED
@@ -98,8 +98,8 @@ Verified across 5 real-world production scenarios:
98
98
 
99
99
  | Tier | Price | Highlights |
100
100
  |---|---|---|
101
- | **Hacker (Open-Core)** | **$0** (Free Forever) | 100k tokens/mo, Local Regex compactor, use own DeepSeek key |
102
- | **Solo Developer** | **$9 / month** | **7-Day Free Trial**, Unlimited Managed DeepSeek-V3, AST Prefetcher, Output Guards |
101
+ | **Hacker (Open-Core)** | **$0** (Free Forever) | 100k tokens/mo, Local Regex compactor, use own DeepSeek / Gemini key |
102
+ | **Solo Developer** | **$9 / month** | **7-Day Free Trial**, Unlimited Managed DeepSeek-Flash / Gemini Flash-Lite, AST Prefetcher, Output Guards |
103
103
  | **Engineering Team** | **$29 / dev / mo** | **7-Day Free Trial**, Team Budget Caps, Centralized Ledger, B2B Docker Proxy |
104
104
  | **Enterprise** | Custom | On-Premise Nuitka C-compiled Docker Proxy, Custom Compliance, SAML SSO |
105
105
 
package/bin/tokensafe.js CHANGED
@@ -41,6 +41,8 @@ Commands:
41
41
  npx tokensafe Interactive setup for Cursor, Antigravity, Claude & Codex
42
42
  npx tokensafe mcp Start Universal MCP Server (for Claude Desktop / Cursor)
43
43
  npx tokensafe proxy [port] Start Transparent Local AI Proxy (default: 8765)
44
+ npx tokensafe graph Analyze repository God Nodes & dependency topology
45
+ npx tokensafe map Generate tokensafe-out/API_SURFACE.md interface snapshot
44
46
  npx tokensafe stats Show token economy savings & performance metrics
45
47
  npx tokensafe benchmark Run zero-loss verification benchmark suite
46
48
  npx tokensafe activate Activate commercial or team seat license key
@@ -72,6 +74,13 @@ switch (command) {
72
74
  scriptArgs = args.slice(1);
73
75
  break;
74
76
 
77
+ case 'graph':
78
+ case 'map':
79
+ case 'skeleton':
80
+ scriptPath = path.join(ROOT_DIR, 'src', 'graph_engine.py');
81
+ scriptArgs = args;
82
+ break;
83
+
75
84
  case 'stats':
76
85
  case 'dashboard':
77
86
  scriptPath = path.join(ROOT_DIR, 'src', 'dashboard.py');
package/package.json CHANGED
@@ -1,6 +1,6 @@
1
1
  {
2
2
  "name": "tokensafe",
3
- "version": "1.0.1",
3
+ "version": "1.1.0",
4
4
  "description": "Autonomous Zero-Loss Token Economy & Context Optimizer for AI Coding Agents",
5
5
  "main": "bin/tokensafe.js",
6
6
  "bin": {
package/plugin.json CHANGED
@@ -1,7 +1,7 @@
1
1
  {
2
2
  "name": "agy-token-optimizer",
3
3
  "displayName": "TokenSafe",
4
- "version": "1.0.1",
4
+ "version": "1.1.0",
5
5
  "description": "Zero-Loss Token Economy & Context Optimizer for AI Agents. Compresses prompts, slices AST/tool outputs, and prevents stuck loops.",
6
6
  "logo": "assets/logo.svg",
7
7
  "author": {
package/src/__init__.py CHANGED
@@ -2,5 +2,5 @@
2
2
  TokenSafe: Autonomous Zero-Loss Token Economy & Context Optimizer for AI Coding Agents.
3
3
  """
4
4
 
5
- __version__ = "1.0.1"
5
+ __version__ = "1.1.0"
6
6
  __author__ = "TokenSafe Team"
@@ -0,0 +1,471 @@
1
+ """
2
+ TokenSafe Graph Engine: Dependency-Aware Graph Analysis & Skeletonization.
3
+ Builds a zero-LLM deterministic code graph (AST-based) for Python, TypeScript, and Go.
4
+ Powers:
5
+ 1. Graph-Guided Skeletonization (target + 1-hop full, 2-hop signature, 0-hop pruned).
6
+ 2. God Nodes Detection (architectural hubs via Degree Centrality).
7
+ 3. API Surface Map Generation (tokensafe-out/API_SURFACE.md).
8
+ 4. Subtree Cache Invalidation & Instant Retrieval via incremental mtime tracking.
9
+ """
10
+
11
+ import os
12
+ import ast
13
+ import re
14
+ import json
15
+ import time
16
+ from typing import Dict, Set, List, Optional, Tuple, Any
17
+
18
+ IGNORED_DIRS = {
19
+ ".git", ".hg", ".svn", "node_modules", "dist", "build", "venv", ".venv",
20
+ "__pycache__", ".idea", ".vscode", "data", "coverage", ".pytest_cache",
21
+ "Library", "Applications", "Pictures", "Movies", "Music", "Downloads",
22
+ ".cache", ".npm", ".bun", ".rustup", ".cargo", ".local", ".system_generated"
23
+ }
24
+
25
+ SUPPORTED_EXTS = {".py", ".ts", ".js", ".tsx", ".jsx", ".go"}
26
+
27
+ class CodeSymbol:
28
+ """Represents a structural code entity (file, class, or function)."""
29
+ __slots__ = ("name", "kind", "file_path", "line_start", "line_end", "signature", "body")
30
+
31
+ def __init__(
32
+ self,
33
+ name: str,
34
+ kind: str,
35
+ file_path: str,
36
+ line_start: int = 0,
37
+ line_end: int = 0,
38
+ signature: str = "",
39
+ body: str = ""
40
+ ):
41
+ self.name = name
42
+ self.kind = kind
43
+ self.file_path = file_path
44
+ self.line_start = line_start
45
+ self.line_end = line_end
46
+ self.signature = signature
47
+ self.body = body
48
+
49
+ def to_dict(self) -> Dict[str, Any]:
50
+ return {
51
+ "name": self.name,
52
+ "kind": self.kind,
53
+ "file_path": self.file_path,
54
+ "line_start": self.line_start,
55
+ "line_end": self.line_end,
56
+ "signature": self.signature
57
+ }
58
+
59
+ class CodeGraph:
60
+ """
61
+ Deterministic In-Memory Knowledge Graph of a software repository.
62
+ Nodes: File and Symbol identifiers.
63
+ Edges: Directed dependencies (CALLS, IMPORTS, CONTAINS).
64
+ """
65
+
66
+ def __init__(self, root_dir: str = "."):
67
+ self.root_dir = os.path.abspath(root_dir)
68
+ self.symbols: Dict[str, CodeSymbol] = {}
69
+ self.file_symbols: Dict[str, List[str]] = {} # file_rel_path -> [symbol_keys]
70
+ self.edges: Dict[str, Set[str]] = {} # caller/importer -> set(callees/imported)
71
+ self.rev_edges: Dict[str, Set[str]] = {} # callee/imported -> set(callers/importers)
72
+ self.file_mtimes: Dict[str, float] = {} # file_rel_path -> mtime
73
+ self.last_scanned: float = 0.0
74
+
75
+ def add_edge(self, source: str, target: str):
76
+ if not source or not target or source == target:
77
+ return
78
+ if source not in self.edges:
79
+ self.edges[source] = set()
80
+ self.edges[source].add(target)
81
+
82
+ if target not in self.rev_edges:
83
+ self.rev_edges[target] = set()
84
+ self.rev_edges[target].add(source)
85
+
86
+ def scan(self, force: bool = False) -> int:
87
+ """
88
+ Fast incremental scan of the workspace.
89
+ Parses modified files only; returns number of files re-indexed.
90
+ """
91
+ reindexed = 0
92
+ current_files = set()
93
+
94
+ for dirpath, dirnames, filenames in os.walk(self.root_dir):
95
+ dirnames[:] = [d for d in dirnames if d not in IGNORED_DIRS and not d.startswith(".")]
96
+ for filename in filenames:
97
+ _, ext = os.path.splitext(filename)
98
+ if ext.lower() in SUPPORTED_EXTS:
99
+ full_path = os.path.join(dirpath, filename)
100
+ rel_path = os.path.relpath(full_path, self.root_dir).replace("\\", "/")
101
+ current_files.add(rel_path)
102
+
103
+ try:
104
+ mtime = os.path.getmtime(full_path)
105
+ except OSError:
106
+ continue
107
+
108
+ if force or rel_path not in self.file_mtimes or self.file_mtimes[rel_path] < mtime:
109
+ self._index_file(full_path, rel_path)
110
+ self.file_mtimes[rel_path] = mtime
111
+ reindexed += 1
112
+
113
+ # Purge removed files
114
+ deleted = set(self.file_mtimes.keys()) - current_files
115
+ for del_path in deleted:
116
+ self._purge_file(del_path)
117
+ self.file_mtimes.pop(del_path, None)
118
+
119
+ self.last_scanned = time.time()
120
+ return reindexed
121
+
122
+ def _purge_file(self, rel_path: str):
123
+ """Removes symbols and edges belonging to a file."""
124
+ sym_keys = self.file_symbols.pop(rel_path, [])
125
+ for k in sym_keys:
126
+ self.symbols.pop(k, None)
127
+ targets = self.edges.pop(k, set())
128
+ for t in targets:
129
+ if t in self.rev_edges:
130
+ self.rev_edges[t].discard(k)
131
+ sources = self.rev_edges.pop(k, set())
132
+ for s in sources:
133
+ if s in self.edges:
134
+ self.edges[s].discard(k)
135
+
136
+ # File-level node
137
+ file_node = f"file:{rel_path}"
138
+ self.edges.pop(file_node, None)
139
+ self.rev_edges.pop(file_node, None)
140
+
141
+ def _index_file(self, full_path: str, rel_path: str):
142
+ """Indexes a single code file using standard AST or regex tokenizers."""
143
+ self._purge_file(rel_path)
144
+ try:
145
+ with open(full_path, "r", encoding="utf-8", errors="replace") as f:
146
+ content = f.read()
147
+ except OSError:
148
+ return
149
+
150
+ file_node = f"file:{rel_path}"
151
+ self.file_symbols[rel_path] = []
152
+
153
+ _, ext = os.path.splitext(full_path)
154
+ ext = ext.lower()
155
+
156
+ if ext == ".py":
157
+ self._parse_python(content, rel_path, file_node)
158
+ elif ext in (".ts", ".js", ".tsx", ".jsx", ".go"):
159
+ self._parse_braces(content, rel_path, file_node, ext)
160
+
161
+ def _parse_python(self, code: str, rel_path: str, file_node: str):
162
+ try:
163
+ tree = ast.parse(code)
164
+ except SyntaxError:
165
+ return
166
+
167
+ lines = code.splitlines()
168
+
169
+ for node in ast.walk(tree):
170
+ if isinstance(node, (ast.Import, ast.ImportFrom)):
171
+ # Record imports
172
+ if isinstance(node, ast.Import):
173
+ for alias in node.names:
174
+ self.add_edge(file_node, f"mod:{alias.name}")
175
+ elif isinstance(node, ast.ImportFrom) and node.module:
176
+ self.add_edge(file_node, f"mod:{node.module}")
177
+ for alias in node.names:
178
+ self.add_edge(file_node, f"sym:{alias.name}")
179
+
180
+ elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
181
+ kind = "class" if isinstance(node, ast.ClassDef) else "function"
182
+ start = getattr(node, "lineno", 1)
183
+ end = getattr(node, "end_lineno", start)
184
+ sym_key = f"{rel_path}::{node.name}"
185
+ sig = lines[start - 1].strip() if 0 < start <= len(lines) else node.name
186
+
187
+ symbol = CodeSymbol(
188
+ name=node.name,
189
+ kind=kind,
190
+ file_path=rel_path,
191
+ line_start=start,
192
+ line_end=end,
193
+ signature=sig
194
+ )
195
+ self.symbols[sym_key] = symbol
196
+ self.file_symbols[rel_path].append(sym_key)
197
+ self.add_edge(file_node, sym_key)
198
+
199
+ # Scan function calls inside body
200
+ for sub in ast.walk(node):
201
+ if isinstance(sub, ast.Call):
202
+ callee_name = None
203
+ if isinstance(sub.func, ast.Name):
204
+ callee_name = sub.func.id
205
+ elif isinstance(sub.func, ast.Attribute):
206
+ callee_name = sub.func.attr
207
+ if callee_name:
208
+ self.add_edge(sym_key, f"call:{callee_name}")
209
+
210
+ def _parse_braces(self, code: str, rel_path: str, file_node: str, ext: str):
211
+ lines = code.splitlines()
212
+
213
+ # Imports
214
+ import_pattern = re.compile(r"""(?:import\s+.*?\s+from\s+['"](.*?)['"]|require\(['"](.*?)['"]\))""")
215
+ for match in import_pattern.finditer(code):
216
+ target_mod = match.group(1) or match.group(2)
217
+ if target_mod:
218
+ self.add_edge(file_node, f"mod:{target_mod}")
219
+
220
+ # Functions & Methods
221
+ fn_pattern = re.compile(r"""(?:export\s+)?(?:async\s+)?(?:function\s+([a-zA-Z0-9_]+)|(?:class|interface)\s+([a-zA-Z0-9_]+)|func\s+(?:\(.*?\)\s+)?([a-zA-Z0-9_]+)|(?:const|let)\s+([a-zA-Z0-9_]+)\s*=\s*(?:async\s*)?\()""")
222
+ for idx, line in enumerate(lines, 1):
223
+ match = fn_pattern.search(line)
224
+ if match:
225
+ name = match.group(1) or match.group(2) or match.group(3) or match.group(4)
226
+ if name:
227
+ kind = "class" if "class" in line or "interface" in line else "function"
228
+ sym_key = f"{rel_path}::{name}"
229
+ symbol = CodeSymbol(
230
+ name=name,
231
+ kind=kind,
232
+ file_path=rel_path,
233
+ line_start=idx,
234
+ line_end=idx,
235
+ signature=line.strip()
236
+ )
237
+ self.symbols[sym_key] = symbol
238
+ self.file_symbols[rel_path].append(sym_key)
239
+ self.add_edge(file_node, sym_key)
240
+
241
+ # -------------------------------------------------------------
242
+ # 1. God Nodes Analysis (Degree Centrality)
243
+ # -------------------------------------------------------------
244
+ def get_god_nodes(self, top_n: int = 5) -> List[Dict[str, Any]]:
245
+ """
246
+ Calculates Degree Centrality across files to locate architectural hubs.
247
+ God Nodes = files with highest incoming (depended upon) and outgoing connections.
248
+ """
249
+ scores: Dict[str, Dict[str, Any]] = {}
250
+
251
+ for rel_path in self.file_symbols.keys():
252
+ file_node = f"file:{rel_path}"
253
+ # Outgoing: What this file imports or contains
254
+ out_degree = len(self.edges.get(file_node, set()))
255
+ # Incoming: Who imports or calls this file or its symbols
256
+ in_degree = len(self.rev_edges.get(file_node, set()))
257
+
258
+ # Aggregate symbol connections
259
+ for sym_key in self.file_symbols.get(rel_path, []):
260
+ sym_name = sym_key.split("::")[-1]
261
+ call_ref = f"call:{sym_name}"
262
+ in_degree += len(self.rev_edges.get(call_ref, set()))
263
+
264
+ total_centrality = in_degree * 2 + out_degree # Incoming imports weighted heavier
265
+ scores[rel_path] = {
266
+ "file": rel_path,
267
+ "centrality": total_centrality,
268
+ "dependents": in_degree,
269
+ "dependencies": out_degree,
270
+ "symbols_count": len(self.file_symbols.get(rel_path, []))
271
+ }
272
+
273
+ ranked = sorted(scores.values(), key=lambda x: x["centrality"], reverse=True)
274
+ return ranked[:top_n]
275
+
276
+ # -------------------------------------------------------------
277
+ # 2. Neighborhood & Graph-Guided Skeletonization
278
+ # -------------------------------------------------------------
279
+ def get_neighborhood(self, target_symbol: str) -> Dict[str, Set[str]]:
280
+ """
281
+ Returns 0-hop (target), 1-hop (callees/callers), and 2-hop related symbols.
282
+ """
283
+ target_keys = [k for k in self.symbols if k.endswith(f"::{target_symbol}") or k == target_symbol]
284
+ one_hop = set()
285
+
286
+ for tk in target_keys:
287
+ # Direct callees
288
+ for callee in self.edges.get(tk, set()):
289
+ if callee.startswith("call:"):
290
+ name = callee[5:]
291
+ for k in self.symbols:
292
+ if k.endswith(f"::{name}"):
293
+ one_hop.add(k)
294
+ elif callee in self.symbols:
295
+ one_hop.add(callee)
296
+ # Direct callers
297
+ for caller in self.rev_edges.get(f"call:{target_symbol}", set()):
298
+ if caller in self.symbols:
299
+ one_hop.add(caller)
300
+
301
+ return {
302
+ "target": set(target_keys),
303
+ "one_hop": one_hop
304
+ }
305
+
306
+ def graph_guided_skeletonize(self, full_path: str, target_symbol: Optional[str] = None) -> str:
307
+ """
308
+ Surgically folds code based on graph connectivity:
309
+ - target_symbol: Retains FULL function body.
310
+ - 1-hop dependencies: Retains FULL body.
311
+ - 2-hop / Unrelated functions in the same file: Stripped to AST signatures (ellipsis / pass).
312
+ Yields up to 90% token reduction without breaking context.
313
+ """
314
+ try:
315
+ with open(full_path, "r", encoding="utf-8", errors="replace") as f:
316
+ code = f.read()
317
+ except OSError:
318
+ return ""
319
+
320
+ if not target_symbol:
321
+ # Fall back to standard skeletonization
322
+ from src.skeletonizer import Skeletonizer
323
+ _, ext = os.path.splitext(full_path)
324
+ return Skeletonizer.skeletonize_code(code, ext)
325
+
326
+ _, ext = os.path.splitext(full_path)
327
+ ext = ext.lower()
328
+
329
+ if ext != ".py":
330
+ # For brace-based languages, fallback to standard skeleton
331
+ from src.skeletonizer import Skeletonizer
332
+ return Skeletonizer.skeletonize_code(code, ext)
333
+
334
+ # Python Graph-Guided AST Pruning
335
+ try:
336
+ tree = ast.parse(code)
337
+ neighborhood = self.get_neighborhood(target_symbol)
338
+ keep_full_names = {target_symbol}
339
+ for k in neighborhood["one_hop"]:
340
+ keep_full_names.add(k.split("::")[-1])
341
+
342
+ for node in ast.walk(tree):
343
+ if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
344
+ if node.name not in keep_full_names:
345
+ doc = ast.get_docstring(node)
346
+ new_body = []
347
+ if doc:
348
+ new_body.append(ast.Expr(value=ast.Constant(value=doc)))
349
+ new_body.append(ast.Expr(value=ast.Constant(value=Ellipsis)))
350
+ node.body = new_body
351
+
352
+ return ast.unparse(tree)
353
+ except Exception:
354
+ return code
355
+
356
+ # -------------------------------------------------------------
357
+ # 3. API Surface Generator (tokensafe-out/API_SURFACE.md)
358
+ # -------------------------------------------------------------
359
+ def generate_api_surface(self, output_path: Optional[str] = None) -> str:
360
+ """
361
+ Generates a high-density, agent-friendly Markdown map of all classes,
362
+ types, and exported interfaces in the project (~1,500 tokens).
363
+ """
364
+ god_nodes = self.get_god_nodes(top_n=5)
365
+ lines = [
366
+ "# TokenSafe Architecture & API Surface",
367
+ f"*Generated locally at {time.strftime('%Y-%m-%d %H:%M:%S')} • Zero LLM Cost*",
368
+ "",
369
+ "## Architectural Hubs (God Nodes)",
370
+ "The foundational modules that connect the largest portion of the project:",
371
+ ""
372
+ ]
373
+
374
+ for g in god_nodes:
375
+ lines.append(f"- **`{g['file']}`**: {g['dependents']} dependents, {g['dependencies']} dependencies ({g['symbols_count']} exported symbols)")
376
+
377
+ lines.extend([
378
+ "",
379
+ "## Interface Surface Map",
380
+ ""
381
+ ])
382
+
383
+ for rel_path, sym_keys in sorted(self.file_symbols.items()):
384
+ if not sym_keys:
385
+ continue
386
+ lines.append(f"### `{rel_path}`")
387
+ for k in sym_keys:
388
+ sym = self.symbols.get(k)
389
+ if sym:
390
+ lines.append(f"- `{sym.kind}` **`{sym.name}`** (line {sym.line_start}): `{sym.signature}`")
391
+ lines.append("")
392
+
393
+ output_content = "\n".join(lines)
394
+
395
+ if output_path is None:
396
+ out_dir = os.path.join(self.root_dir, "tokensafe-out")
397
+ os.makedirs(out_dir, exist_ok=True)
398
+ output_path = os.path.join(out_dir, "API_SURFACE.md")
399
+
400
+ try:
401
+ with open(output_path, "w", encoding="utf-8") as f:
402
+ f.write(output_content)
403
+ except OSError:
404
+ pass
405
+
406
+ return output_content
407
+
408
+ def format_topology_banner(self) -> str:
409
+ """Creates a compact 1-line prompt injection banner for coding agents."""
410
+ god_nodes = self.get_god_nodes(top_n=4)
411
+ if not god_nodes:
412
+ return ""
413
+ items = [f"{g['file']} ({g['dependents']} deps)" for g in god_nodes]
414
+ return f"[TokenSafe Architecture Map | Core Hubs: {', '.join(items)}]"
415
+
416
+
417
+ class GraphEngine:
418
+ """Singleton Manager for automatic background graph execution."""
419
+
420
+ _instance: Optional[CodeGraph] = None
421
+
422
+ @classmethod
423
+ def get_instance(cls, root_dir: str = ".") -> CodeGraph:
424
+ if cls._instance is None or cls._instance.root_dir != os.path.abspath(root_dir):
425
+ cls._instance = CodeGraph(root_dir)
426
+ cls._instance.scan()
427
+ else:
428
+ # Incremental scan if older than 3 seconds
429
+ if time.time() - cls._instance.last_scanned > 3.0:
430
+ cls._instance.scan()
431
+ return cls._instance
432
+
433
+ def main():
434
+ import sys
435
+ args = sys.argv[1:]
436
+ cmd = args[0] if args else "graph"
437
+
438
+ target_dir = "."
439
+ graph = GraphEngine.get_instance(target_dir)
440
+
441
+ if cmd in ("graph", "nodes", "god-nodes"):
442
+ print(f"\nšŸ” TokenSafe Graph Engine — Knowledge Graph & Centrality Analysis")
443
+ print(f" Indexed: {len(graph.file_symbols)} code files, {len(graph.symbols)} symbols\n")
444
+ god_nodes = graph.get_god_nodes(top_n=7)
445
+ print("⚔ Architectural Hubs (God Nodes):")
446
+ for idx, g in enumerate(god_nodes, 1):
447
+ print(f" {idx}. {g['file']}")
448
+ print(f" Degree Centrality: {g['centrality']} | Dependents: {g['dependents']} | Dependencies: {g['dependencies']} | Symbols: {g['symbols_count']}")
449
+ print(f"\nšŸ’” Tip: Run 'tokensafe map' to generate tokensafe-out/API_SURFACE.md\n")
450
+
451
+ elif cmd in ("map", "surface"):
452
+ out_file = os.path.join(target_dir, "tokensafe-out", "API_SURFACE.md")
453
+ graph.generate_api_surface(out_file)
454
+ print(f"\nāœ“ Generated API Surface Map: {out_file} ({os.path.getsize(out_file):,} bytes)")
455
+ print(f" AI agents can read this file in ~1.5k tokens instead of reading the entire repo.\n")
456
+
457
+ elif cmd == "skeleton":
458
+ if len(args) < 2:
459
+ print("Usage: tokensafe skeleton <file_path> [target_symbol]")
460
+ sys.exit(1)
461
+ file_path = args[1]
462
+ target_symbol = args[2] if len(args) > 2 else None
463
+ res = graph.graph_guided_skeletonize(file_path, target_symbol)
464
+ print(res)
465
+
466
+ else:
467
+ print(f"Unknown command: {cmd}. Available: graph, map, skeleton")
468
+
469
+ if __name__ == "__main__":
470
+ main()
471
+
package/src/proxy.py CHANGED
@@ -21,6 +21,7 @@ from src.metrics import MetricsTracker
21
21
  from src.dashboard import Dashboard
22
22
  from src.loop_detector import LoopDetector
23
23
  from src.secret_redactor import SecretRedactor
24
+ from src.graph_engine import GraphEngine
24
25
 
25
26
  DEFAULT_PORT = 8765
26
27
  DEFAULT_UPSTREAM = "https://api.deepseek.com/v1"
@@ -86,6 +87,7 @@ class TokenSafeProxyHandler(BaseHTTPRequestHandler):
86
87
  print(f"[TokenSafe] Scrubbed {secrets_scrubbed} sensitive secret(s) from payload before transmission.")
87
88
 
88
89
  # 2. Ingress Optimization: Compact last user message
90
+ target_symbol = None
89
91
  if messages:
90
92
  for idx in reversed(range(len(messages))):
91
93
  msg = messages[idx]
@@ -94,9 +96,22 @@ class TokenSafeProxyHandler(BaseHTTPRequestHandler):
94
96
  compaction = PromptCompactor.compact(raw_user, enable_deepseek=False)
95
97
  if compaction["saved_tokens"] > 5:
96
98
  msg["content"] = compaction["compact_text"]
99
+ target_symbol = compaction.get("symbol")
97
100
  break
98
101
 
99
- # 3. Forward to Upstream
102
+ # 3. Automatic Architecture Awareness (God Nodes Injection)
103
+ if messages:
104
+ try:
105
+ graph = GraphEngine.get_instance(os.getcwd())
106
+ banner = graph.format_topology_banner()
107
+ first_msg = messages[0]
108
+ if banner and isinstance(first_msg.get("content"), str) and "[TokenSafe Architecture" not in first_msg["content"]:
109
+ if first_msg.get("role") == "system":
110
+ first_msg["content"] = f"{banner}\n{first_msg['content']}"
111
+ except Exception:
112
+ pass
113
+
114
+ # 4. Forward to Upstream
100
115
  auth_header = self.headers.get("Authorization", "")
101
116
  # Check if local config has deepseek key
102
117
  config_path = os.path.join(BASE_DIR, "data", "config.json")
@@ -140,11 +155,24 @@ class TokenSafeProxyHandler(BaseHTTPRequestHandler):
140
155
  self.wfile.write(json.dumps({"error": f"Upstream proxy error: {str(e)}"}).encode("utf-8"))
141
156
 
142
157
  def run_proxy(port: int = DEFAULT_PORT):
158
+ # Auto-initialize Knowledge Graph & API Surface on startup
159
+ hubs_summary = "active"
160
+ try:
161
+ graph = GraphEngine.get_instance(os.getcwd())
162
+ graph.generate_api_surface()
163
+ god_nodes = graph.get_god_nodes(3)
164
+ if god_nodes:
165
+ hubs_summary = ", ".join([f"{g['file']} ({g['dependents']} deps)" for g in god_nodes])
166
+ except Exception as e:
167
+ hubs_summary = f"initialized ({e})"
168
+
143
169
  server_address = ("127.0.0.1", port)
144
170
  httpd = HTTPServer(server_address, TokenSafeProxyHandler)
145
171
  print(f"\n⚔ TokenSafe Local AI Proxy running on: http://127.0.0.1:{port}")
146
172
  print(f" Dashboard: http://127.0.0.1:{port}/dashboard")
147
173
  print(f" OpenAI-compatible URL: http://127.0.0.1:{port}/v1")
174
+ print(f" Auto-Graph: {hubs_summary}")
175
+ print(f" API Surface: tokensafe-out/API_SURFACE.md")
148
176
  try:
149
177
  httpd.serve_forever()
150
178
  except KeyboardInterrupt: