gemcode 0.2.2__py3-none-any.whl

gemcode/model_routing.py

@@ -0,0 +1,179 @@
+"""
+Model routing for GemCode.
+
+Goal: match Claude-style "multi modes" where users can select a model mode
+explicitly, or GemCode can choose a best-fit model automatically (no extra
+model call; heuristic only).
+"""
+
+from __future__ import annotations
+
+import re
+from typing import Literal
+
+from gemcode.config import GemCodeConfig
+
+ModelMode = Literal["auto", "fast", "balanced", "quality"]
+FamilyMode = Literal["auto", "primary", "alt"]
+
+
+def _contains_any(haystack: str, needles: list[str]) -> bool:
+  h = haystack.lower()
+  return any(n in h for n in needles)
+
+
+def pick_effective_model(cfg: GemCodeConfig, prompt: str) -> str:
+  """
+  Returns the effective model id to use for this run.
+
+  Heuristic rules (cheap, deterministic):
+  - quality for architecture/design/refactor/complex trade-offs
+  - fast for small edits, tool-driven tasks, or quick debugging
+  - balanced as the default otherwise
+  """
+  mode_norm = (getattr(cfg, "model_mode", "fast") or "fast").lower()
+
+  # If the user explicitly picked a model id, honor it.
+  if getattr(cfg, "model_overridden", False):
+    return cfg.model
+
+  # Optional deep research routing: when user asks explicitly (flag) or when
+  # `model_mode=auto` and the prompt requests research-like output.
+  deep_research_triggers = [
+    "deep research",
+    "deep-dive",
+    "research",
+    "sources",
+    "citations",
+    "grounded",
+    "investigate",
+    "literature",
+    "benchmark",
+  ]
+  prompt_norm = re.sub(r"\s+", " ", prompt or "").strip().lower()
+  if getattr(cfg, "enable_deep_research", False):
+    return getattr(cfg, "model_deep_research", None) or cfg.model
+  if mode_norm == "auto" and _contains_any(prompt_norm, deep_research_triggers):
+    return getattr(cfg, "model_deep_research", None) or cfg.model
+
+  # Capability precedence: computer-use model selection.
+  # (Deep research already handled above.)
+  if getattr(cfg, "enable_audio", False):
+    return getattr(cfg, "model_audio_live", None) or cfg.model
+  if getattr(cfg, "enable_computer_use", False):
+    return getattr(cfg, "model_computer_use", None) or cfg.model
+
+  primary_fast = cfg.model
+  primary_quality = getattr(cfg, "model_quality", None) or primary_fast
+  primary_balanced = getattr(cfg, "model_balanced", None) or primary_fast
+
+  alt_fast = getattr(cfg, "model_alt", None) or primary_fast
+  alt_quality = getattr(cfg, "model_alt_quality", None) or primary_quality
+  alt_balanced = getattr(cfg, "model_alt_balanced", None) or primary_balanced
+
+  if mode_norm not in ("auto", "fast", "balanced", "quality"):
+    return primary_fast
+
+  def decide_base_mode() -> Literal["fast", "balanced", "quality"]:
+    if mode_norm == "fast":
+      return "fast"
+    if mode_norm == "balanced":
+      return "balanced"
+    if mode_norm == "quality":
+      return "quality"
+
+    # auto mode: choose base mode using prompt heuristics.
+    p_norm = re.sub(r"\s+", " ", prompt or "").strip()
+    plen = len(p_norm)
+
+    quality_triggers = [
+      "architecture",
+      "design",
+      "system",
+      "refactor",
+      "trade",
+      "complex",
+      "scal",
+      "performance",
+      "profil",
+      "migration",
+      "schema",
+      "how would you",
+      "deep dive",
+    ]
+    fast_triggers = [
+      "fix",
+      "bug",
+      "error",
+      "tests",
+      "pytest",
+      "debug",
+      "failing",
+      "lint",
+      "format",
+      "quick",
+      "small change",
+    ]
+
+    if (
+      plen > 2_000
+      and _contains_any(p_norm, ["design", "architecture", "refactor", "trade"])
+    ):
+      return "quality"
+
+    if _contains_any(p_norm, quality_triggers):
+      return "quality"
+
+    if _contains_any(p_norm, fast_triggers):
+      return "fast"
+
+    # If prompt is long but not explicitly complex, balanced tends to be safer.
+    if plen > 4_000:
+      return "balanced"
+
+    return "balanced"
+
+  base_mode = decide_base_mode()
+
+  # Decide model family (primary vs 2.5-alt).
+  fam = (getattr(cfg, "model_family_mode", "auto") or "auto").lower()
+  if fam not in ("auto", "primary", "alt"):
+    fam = "auto"
+
+  p_norm2 = re.sub(r"\s+", " ", prompt or "").strip()
+  # Reuse quality triggers: complex prompts get primary (3.x); simpler prompts
+  # prefer alt (2.5) by default in `auto` family mode.
+  complex_triggers = [
+    "architecture",
+    "design",
+    "system",
+    "refactor",
+    "trade",
+    "complex",
+    "performance",
+    "migration",
+    "schema",
+    "deep dive",
+  ]
+
+  choose_primary: bool
+  if fam == "primary":
+    choose_primary = True
+  elif fam == "alt":
+    choose_primary = False
+  else:
+    choose_primary = _contains_any(p_norm2, complex_triggers)
+
+  if choose_primary:
+    if base_mode == "fast":
+      return primary_fast
+    if base_mode == "balanced":
+      return primary_balanced
+    return primary_quality
+
+  if base_mode == "fast":
+    return alt_fast
+  if base_mode == "balanced":
+    return alt_balanced
+  return alt_quality
+

gemcode/paths.py

@@ -0,0 +1,29 @@
+"""Safe path resolution under a project root."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+
+class PathEscapeError(ValueError):
+  """Resolved path would leave the project sandbox."""
+
+
+def resolve_under_root(project_root: Path, rel: str) -> Path:
+  """
+  Resolve a user-relative path against project_root.
+
+  Rejects absolute paths that escape the root (symlink traversal is still a
+  concern for production; MVP uses realpath check).
+  """
+  root = project_root.resolve()
+  raw = Path(rel)
+  if raw.is_absolute():
+    candidate = raw.resolve()
+  else:
+    candidate = (root / raw).resolve()
+  try:
+    candidate.relative_to(root)
+  except ValueError as e:
+    raise PathEscapeError(f"Path outside project root: {rel!r}") from e
+  return candidate

gemcode/permissions.py

@@ -0,0 +1,5 @@
+"""Backward-compatible re-exports. Prefer `gemcode.callbacks` for new code."""
+
+from gemcode.callbacks import make_before_tool_callback
+
+__all__ = ["make_before_tool_callback"]

gemcode/plugins/terminal_hooks_plugin.py

@@ -0,0 +1,168 @@
+"""
+ADK plugin: complements Claude-like stopHooks with GemCode terminal reasons,
+optional memory ingestion, and post-turn hook execution.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import os
+
+from google.adk.plugins.base_plugin import BasePlugin
+
+from google.adk.models.google_llm import Gemini
+from google.adk.models.llm_request import LlmRequest
+from google.genai import types
+
+from gemcode.config import GemCodeConfig
+from gemcode.query.stop_hooks import run_post_turn_hooks
+from gemcode.audit import append_audit
+from gemcode.prompt_suggestions import build_prompt_suggestion
+
+
+class GemCodeTerminalHooksPlugin(BasePlugin):
+  def __init__(self, cfg: GemCodeConfig):
+    super().__init__(name="gemcode_terminal_hooks")
+    self.cfg = cfg
+
+  def _use_interactions_for_prompt_suggestions(self) -> bool:
+    v = os.environ.get("GEMCODE_PROMPT_SUGGESTIONS_USE_INTERACTIONS", "1")
+    return v.lower() in ("1", "true", "yes", "on")
+
+  def _find_previous_interaction_id(
+    self, *, callback_context: Any, agent_name: str
+  ) -> str | None:
+    # Interactions chaining uses `previous_interaction_id` extracted from the
+    # most recent model response event for the same agent.
+    try:
+      events = callback_context.session.events or []
+    except Exception:
+      return None
+    for ev in reversed(events):
+      if getattr(ev, "author", None) == agent_name and getattr(
+        ev, "interaction_id", None
+      ):
+        return ev.interaction_id
+    return None
+
+  async def _suggest_via_interactions(
+    self,
+    *,
+    terminal_reason: str,
+    callback_context: Any,
+    agent: Any,
+    heuristic: str,
+  ) -> str | None:
+    if not self._use_interactions_for_prompt_suggestions():
+      return None
+
+    try:
+      previous_interaction_id = self._find_previous_interaction_id(
+        callback_context=callback_context,
+        agent_name=getattr(agent, "name", "gemcode"),
+      )
+
+      prompt = (
+        "You are GemCode. Provide the best next-step guidance for the user. "
+        "Terminal reason: {reason}. "
+        "Write a single short sentence (<= 220 chars). "
+        "If it involves policy/actions, reference exact flags like `--yes` or `--session`.\n\n"
+        "Heuristic suggestion (may be imperfect): {heuristic}"
+      ).format(reason=terminal_reason, heuristic=heuristic)
+
+      llm = Gemini(model=self.cfg.model, use_interactions_api=True)
+      req = LlmRequest(
+        model=self.cfg.model,
+        contents=[
+          types.Content(
+            role="user",
+            parts=[types.Part(text=prompt)],
+          )
+        ],
+        config=types.GenerateContentConfig(),
+      )
+      if previous_interaction_id:
+        req.previous_interaction_id = previous_interaction_id
+
+      async for resp in llm.generate_content_async(req, stream=False):
+        if resp.content and resp.content.parts:
+          texts = [
+            getattr(p, "text", None)
+            for p in resp.content.parts
+            if getattr(p, "text", None)
+          ]
+          if texts:
+            out = "".join(texts).strip()
+            if out:
+              return out
+      return None
+    except Exception as e:
+      append_audit(
+        self.cfg.project_root,
+        {
+          "phase": "prompt_suggestion_interactions",
+          "ok": False,
+          "error": str(e),
+          "terminal_reason": terminal_reason,
+        },
+      )
+      return None
+
+  async def after_agent_callback(self, *, agent: Any, callback_context: Any):
+    # callback_context is an ADK Context (mutable state + helper methods).
+    state = callback_context.state
+    terminal_reason = state.get("gemcode:terminal_reason", None)
+    if not terminal_reason:
+      terminal_reason = "completed"
+
+    append_audit(self.cfg.project_root, {"phase": "terminal", "reason": terminal_reason})
+
+    heuristic = build_prompt_suggestion(
+      self.cfg, terminal_reason=terminal_reason
+    )
+    if heuristic:
+      suggestion = heuristic
+      suggestion_via_interactions = await self._suggest_via_interactions(
+        terminal_reason=terminal_reason,
+        callback_context=callback_context,
+        agent=agent,
+        heuristic=heuristic,
+      )
+      if suggestion_via_interactions:
+        suggestion = suggestion_via_interactions
+
+      append_audit(
+        self.cfg.project_root,
+        {
+          "phase": "prompt_suggestion",
+          "terminal_reason": terminal_reason,
+          "suggestion": suggestion,
+        },
+      )
+
+    if getattr(self.cfg, "enable_memory", False):
+      try:
+        await callback_context.add_session_to_memory()
+        append_audit(self.cfg.project_root, {"phase": "memory_ingest", "ok": True})
+      except Exception as e:
+        append_audit(
+          self.cfg.project_root,
+          {"phase": "memory_ingest", "ok": False, "error": str(e)},
+        )
+
+    # Execute stopHooks-like script hook at the end of the invocation.
+    try:
+      run_post_turn_hooks(
+        self.cfg,
+        session_id=callback_context.session.id,
+        user_id=callback_context.user_id,
+      )
+    except Exception as e:
+      append_audit(
+        self.cfg.project_root,
+        {"phase": "post_turn_hook", "ok": False, "error": str(e)},
+      )
+
+    return None
+

gemcode/plugins/tool_recovery_plugin.py

@@ -0,0 +1,135 @@
+"""
+Claude-like recovery-loop for GemCode tool failures.
+
+We complement ADK's `ReflectAndRetryToolPlugin` by treating our tool result
+dicts like `{"error": "...", "error_kind": "..."}` as retryable tool failures
+so the model gets reflection guidance and can try a corrected approach.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Any, Optional
+
+from google.adk.plugins.reflect_retry_tool_plugin import (
+    ReflectAndRetryToolPlugin,
+    TrackingScope,
+)
+from google.adk.tools.base_tool import BaseTool
+from google.adk.tools.tool_context import ToolContext
+
+from gemcode.audit import append_audit
+from gemcode.config import GemCodeConfig
+
+
+_STATE_FAILURE_KEY = "gemcode:consecutive_tool_failures"
+_TERMINAL_REASON_KEY = "gemcode:terminal_reason"
+
+_ERROR_KIND_PERMISSION_DENIED = "permission_denied"
+_ERROR_KIND_PERMISSION_BLOCK = "permission_block"
+_ERROR_KIND_CIRCUIT_BREAKER = "circuit_breaker"
+
+
+class GemCodeReflectAndRetryToolPlugin(ReflectAndRetryToolPlugin):
+  """Retry tool failures with reflection guidance (Claude recovery-loop)."""
+
+  def __init__(self, cfg: GemCodeConfig):
+    self.cfg = cfg
+
+    enabled = os.environ.get("GEMCODE_ENABLE_TOOL_RECOVERY_RETRY", "1").lower()
+    if enabled not in ("1", "true", "yes", "on"):
+      # Still construct; we just set max_retries=0 so it becomes inert.
+      max_retries = 0
+    else:
+      max_retries = int(os.environ.get("GEMCODE_TOOL_REFLECT_MAX_RETRIES", "1"))
+
+    super().__init__(
+        max_retries=max_retries,
+        throw_exception_if_retry_exceeded=False,
+        tracking_scope=TrackingScope.INVOCATION,
+    )
+
+  async def extract_error_from_result(
+    self,
+    *,
+    tool: BaseTool,
+    tool_args: dict[str, Any],
+    tool_context: ToolContext,
+    result: Any,
+  ) -> Optional[Exception]:
+    """
+    Treat `{ "error": ... }` tool results as retryable failures.
+
+    Important: skip policy rejections (permission denials / circuit breaker)
+    so we don't waste retries on user-actionable gating.
+    """
+    if not isinstance(result, dict):
+      return None
+    if "error" not in result:
+      return None
+
+    err_kind = result.get("error_kind")
+    if err_kind in (
+      _ERROR_KIND_PERMISSION_DENIED,
+      _ERROR_KIND_PERMISSION_BLOCK,
+      _ERROR_KIND_CIRCUIT_BREAKER,
+    ):
+      return None
+
+    # Update GemCode streak/terminal state since canonical agent callbacks
+    # are likely short-circuited when this plugin returns a reflection.
+    try:
+      st = tool_context.state
+      st[_STATE_FAILURE_KEY] = st.get(_STATE_FAILURE_KEY, 0) + 1
+      if not st.get(_TERMINAL_REASON_KEY):
+        st[_TERMINAL_REASON_KEY] = "tool_retryable_error"
+    except Exception:
+      pass
+
+    err = result.get("error")
+    err_text = err if isinstance(err, str) else str(err)
+    append_audit(
+        self.cfg.project_root,
+        {
+          "phase": "tool_recovery_retry",
+          "tool": tool.name,
+          "error_kind": err_kind,
+          "error": err_text,
+        },
+    )
+
+    return Exception(err_text)
+
+  async def on_tool_error_callback(
+    self,
+    *,
+    tool: BaseTool,
+    tool_args: dict[str, Any],
+    tool_context: ToolContext,
+    error: Exception,
+  ) -> Optional[dict[str, Any]]:
+    """Ensure GemCode streak/terminal state is updated on exceptions."""
+    try:
+      st = tool_context.state
+      st[_STATE_FAILURE_KEY] = st.get(_STATE_FAILURE_KEY, 0) + 1
+      if not st.get(_TERMINAL_REASON_KEY):
+        st[_TERMINAL_REASON_KEY] = "tool_exception"
+    except Exception:
+      pass
+
+    append_audit(
+        self.cfg.project_root,
+        {
+          "phase": "tool_recovery_exception",
+          "tool": tool.name,
+          "error": f"{type(error).__name__}: {error}",
+        },
+    )
+
+    return await super().on_tool_error_callback(
+      tool=tool,
+      tool_args=tool_args,
+      tool_context=tool_context,
+      error=error,
+    )
+

gemcode/prompt_suggestions.py

@@ -0,0 +1,80 @@
+"""
+Heuristic next-step guidance (Claude stopHooks-style).
+
+We can't perfectly replicate Claude's UI-level "next suggestion job" without
+extra model calls, but we can produce reliable, deterministic guidance from
+GemCode's recorded `terminal_reason`.
+"""
+
+from __future__ import annotations
+
+import os
+
+from gemcode.config import GemCodeConfig
+
+
+def _truthy_env(name: str, *, default: bool = False) -> bool:
+  v = os.environ.get(name)
+  if v is None:
+    return default
+  return v.lower() in ("1", "true", "yes", "on")
+
+
+def build_prompt_suggestion(
+  cfg: GemCodeConfig, *, terminal_reason: str
+) -> str | None:
+  if not _truthy_env("GEMCODE_ENABLE_PROMPT_SUGGESTIONS", default=True):
+    return None
+
+  r = terminal_reason
+  if r in ("completed", ""):
+    return None
+
+  if r == "permission_denied":
+    return (
+      "Some actions were blocked by policy. Re-run with `--yes` (or add the "
+      "needed command to `GEMCODE_ALLOW_COMMANDS` when in strict mode), then "
+      "try again with the same request."
+    )
+
+  if r == "tool_circuit_breaker":
+    return (
+      "Tool execution is being halted by the circuit breaker. Start a new "
+      "session (`--session <new_id>`), then either fix the failing tool "
+      "inputs or increase `GEMCODE_MAX_CONSECUTIVE_TOOL_FAILURES`."
+    )
+
+  if r == "session_token_limit":
+    return (
+      "This session exceeded the token ceiling. Start a new session (new "
+      "`--session` id) or raise `GEMCODE_MAX_SESSION_TOKENS`, then re-run "
+      "the request."
+    )
+
+  if r == "token_budget_stop":
+    return (
+      "Per-turn token budget was exhausted. Re-run the request (or split it "
+      "into smaller steps). If you want more room, increase "
+      "`GEMCODE_TOKEN_BUDGET` or reduce the prompt/context."
+    )
+
+  if r in ("tool_exception", "tool_retryable_error"):
+    return (
+      "A tool raised an exception. Check `.gemcode/audit.log` for the tool "
+      "error details, then re-run with corrected inputs or fewer/shorter "
+      "arguments."
+    )
+
+  if r == "model_error":
+    return (
+      "The model call failed. Try again, reduce prompt size, or switch to a "
+      "different model via `--model`."
+    )
+
+  # Generic fallback.
+  return (
+    f"The run ended with terminal reason `{terminal_reason}`. Check "
+    f"`.gemcode/audit.log`, then re-run with a narrower request or after "
+    f"adjusting limits/policy flags."
+  )
+

gemcode/query/__init__.py

@@ -0,0 +1,36 @@
+"""
+Query-layer types and helpers (clean-room analogue of Claude Code `src/query/*`).
+
+- `transitions` — terminal vs continue reasons for the model↔tool loop.
+- `config` — immutable gate snapshot per run (env/session).
+- `token_budget` — continuation/stop decisions vs a per-turn token budget.
+- `deps` — injectable dependencies for tests.
+- `stop_hooks` — optional post-turn subprocess hooks.
+- `engine` — `GemCodeQueryEngine` facade (outer session + submit message).
+
+The ADK `Runner` still executes the inner loop; these modules document parity and
+host logic that maps to `query.ts` / `QueryEngine.ts` responsibilities.
+"""
+
+from gemcode.query.config import QueryGates, build_query_gates
+from gemcode.query.token_budget import (
+  BudgetTracker,
+  TokenBudgetDecision,
+  check_token_budget,
+  get_budget_continuation_message,
+)
+from gemcode.query.transitions import Continue, Terminal
+
+# Note: import `GemCodeQueryEngine` from `gemcode.query.engine` to avoid import cycles
+# (engine pulls session_runtime → agent → callbacks → query.token_budget).
+
+__all__ = [
+  "BudgetTracker",
+  "Continue",
+  "QueryGates",
+  "Terminal",
+  "TokenBudgetDecision",
+  "build_query_gates",
+  "check_token_budget",
+  "get_budget_continuation_message",
+]

gemcode/query/config.py

@@ -0,0 +1,35 @@
+"""
+Immutable gate snapshot at query entry (cf. `query/config.ts`).
+
+Feature flags in Claude use `feature()` for tree-shaking; we use env + this struct.
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+
+
+def _truthy(name: str, default: bool = False) -> bool:
+  v = os.environ.get(name)
+  if v is None:
+    return default
+  return v.lower() in ("1", "true", "yes", "on")
+
+
+@dataclass(frozen=True)
+class QueryGates:
+  """Runtime gates (env), snapshotted once per invocation."""
+
+  emit_tool_use_summaries: bool
+  fast_mode_enabled: bool
+  streaming_tool_execution: bool
+
+
+def build_query_gates() -> QueryGates:
+  """Snapshot env-driven gates (no network)."""
+  return QueryGates(
+      emit_tool_use_summaries=_truthy("GEMCODE_EMIT_TOOL_USE_SUMMARIES"),
+      fast_mode_enabled=not _truthy("GEMCODE_DISABLE_FAST_MODE"),
+      streaming_tool_execution=_truthy("GEMCODE_STREAMING_TOOL_EXEC", default=True),
+  )

gemcode/query/deps.py

@@ -0,0 +1,20 @@
+"""
+Injectable dependencies (cf. claude-code `query/deps.ts`).
+
+Kept minimal: tests can replace `uuid` without patching modules.
+"""
+
+from __future__ import annotations
+
+import uuid as uuid_mod
+from dataclasses import dataclass
+from typing import Callable
+
+
+@dataclass
+class QueryDeps:
+  uuid: Callable[[], str]
+
+
+def production_deps() -> QueryDeps:
+  return QueryDeps(uuid=lambda: str(uuid_mod.uuid4()))