tweek 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

tweek/mcp/approval.py

@@ -73,7 +73,7 @@ class ApprovalRequest:
             return False
         try:
             ts = datetime.fromisoformat(self.timestamp)
-            elapsed = (datetime.utcnow() - ts).total_seconds()
+            elapsed = (datetime.now(tz=None) - ts).total_seconds()
             return elapsed >= self.timeout_seconds
         except (ValueError, TypeError):
             return False
@@ -83,7 +83,7 @@ class ApprovalRequest:
         """Seconds remaining before timeout. Returns 0 if expired."""
         try:
             ts = datetime.fromisoformat(self.timestamp)
-            elapsed = (datetime.utcnow() - ts).total_seconds()
+            elapsed = (datetime.now(tz=None) - ts).total_seconds()
             remaining = self.timeout_seconds - elapsed
             return max(0.0, remaining)
         except (ValueError, TypeError):
@@ -261,7 +261,11 @@ class ApprovalQueue:
         return [self._row_to_request(row) for row in rows]
 
     def get_request(self, request_id: str) -> Optional[ApprovalRequest]:
-        """Get a specific approval request by ID (supports short IDs)."""
+        """Get a specific approval request by ID (supports short IDs).
+
+        Raises:
+            ValueError: If short ID matches multiple requests (ambiguous).
+        """
         with self._get_connection() as conn:
             # Try exact match first
             row = conn.execute(
@@ -271,10 +275,18 @@ class ApprovalQueue:
 
             # If not found, try prefix match (short ID)
             if row is None and len(request_id) < 36:
-                row = conn.execute(
+                rows = conn.execute(
                     "SELECT * FROM approval_requests WHERE id LIKE ?",
                     (f"{request_id}%",),
-                ).fetchone()
+                ).fetchall()
+                if len(rows) == 1:
+                    row = rows[0]
+                elif len(rows) > 1:
+                    ids = [r["id"][:12] for r in rows]
+                    raise ValueError(
+                        f"Ambiguous short ID '{request_id}' matches {len(rows)} requests: "
+                        f"{', '.join(ids)}. Use a longer prefix."
+                    )
 
         return self._row_to_request(row) if row else None
 
@@ -300,17 +312,28 @@ class ApprovalQueue:
         if status not in (ApprovalStatus.APPROVED, ApprovalStatus.DENIED, ApprovalStatus.EXPIRED):
             raise ValueError(f"Invalid decision status: {status}")
 
-        # Resolve short IDs
-        request = self.get_request(request_id)
-        if request is None:
-            return False
-        if request.status != ApprovalStatus.PENDING:
-            return False
-
-        full_id = request.id
-
         def _do_decide():
             with self._get_connection() as conn:
+                # Atomic check-and-update: resolve short IDs and update in a
+                # single connection to eliminate the TOCTOU race condition.
+                # For short IDs, do the LIKE lookup inside the same transaction.
+                if len(request_id) < 36:
+                    rows = conn.execute(
+                        "SELECT id FROM approval_requests WHERE id LIKE ? AND status = 'pending'",
+                        (f"{request_id}%",),
+                    ).fetchall()
+                    if len(rows) == 0:
+                        return False
+                    if len(rows) > 1:
+                        ids = [r["id"][:12] for r in rows]
+                        raise ValueError(
+                            f"Ambiguous short ID '{request_id}' matches {len(rows)} requests: "
+                            f"{', '.join(ids)}. Use a longer prefix."
+                        )
+                    full_id = rows[0]["id"]
+                else:
+                    full_id = request_id
+
                 cursor = conn.execute(
                     """
                     UPDATE approval_requests
@@ -452,5 +475,23 @@ class ApprovalQueue:
             redactor = LogRedactor(enabled=True)
             return redactor.redact_dict(arguments)
         except ImportError:
-            # If logging module unavailable, store as-is
-            return arguments
+            # If logging module unavailable, do basic redaction to avoid storing raw secrets
+            return self._basic_redact(arguments)
+
+    @staticmethod
+    def _basic_redact(arguments: Dict[str, Any]) -> Dict[str, Any]:
+        """Fallback redaction when LogRedactor is unavailable."""
+        import re
+        sensitive_keys = re.compile(
+            r'(?i)(password|secret|token|key|credential|auth|bearer|api.?key)', re.IGNORECASE
+        )
+        redacted = {}
+        for k, v in arguments.items():
+            if sensitive_keys.search(k):
+                redacted[k] = "***REDACTED***"
+            elif isinstance(v, str) and len(v) > 50:
+                # Long string values may contain secrets — truncate
+                redacted[k] = v[:8] + "***"
+            else:
+                redacted[k] = v
+        return redacted

tweek/mcp/proxy.py

@@ -12,6 +12,7 @@ Architecture:
 Usage:
     tweek mcp proxy       # Start proxy on stdio transport
 """
+from __future__ import annotations
 
 import asyncio
 import json
@@ -470,6 +471,23 @@ class TweekMCPProxy:
                 text=json.dumps({"result": "empty response from upstream"}),
             )]
 
+        # Screen output for leaked credentials or sensitive data
+        try:
+            from tweek.mcp.screening import run_output_scan
+            combined_text = "\n".join(tc.text for tc in text_contents)
+            scan_result = run_output_scan(combined_text)
+            if scan_result.get("blocked"):
+                reason = scan_result.get("reason", "Output blocked by security screening")
+                return [TextContent(
+                    type="text",
+                    text=json.dumps({
+                        "error": "Output blocked by Tweek security screening",
+                        "reason": reason,
+                    }),
+                )]
+        except Exception:
+            pass  # Output scanning errors should not block the response
+
         return text_contents
 
     def _build_context(

tweek/mcp/screening.py

@@ -122,13 +122,13 @@ def run_mcp_screening(context: ScreeningContext) -> Dict[str, Any]:
         }
 
     except ImportError as e:
-        logger.warning(f"Screening modules not available: {e}")
-        # Fail open with warning if screening not available
+        logger.error(f"Screening modules not available: {e}")
+        # Fail closed: missing security modules should block, not bypass
         return {
-            "allowed": True,
-            "blocked": False,
+            "allowed": False,
+            "blocked": True,
             "should_prompt": False,
-            "reason": f"Warning: screening unavailable ({e})",
+            "reason": f"Screening unavailable (missing modules: {e}). Blocking for safety.",
             "findings": [],
         }
     except Exception as e:

tweek/mcp/server.py

@@ -14,6 +14,7 @@ proxy mode: tweek mcp proxy
 Usage:
     tweek mcp serve       # stdio mode (desktop clients)
 """
+from __future__ import annotations
 
 import json
 import logging
@@ -233,7 +234,9 @@ class TweekMCPServer:
             })
 
         except Exception as e:
-            return json.dumps({"error": str(e)})
+            # Don't leak internal details across trust boundary
+            logger.error(f"Vault operation failed: {e}")
+            return json.dumps({"error": "Vault operation failed"})
 
     async def _handle_status(self, arguments: Dict[str, Any]) -> str:
         """Handle tweek_status tool call."""

tweek/memory/__init__.py

@@ -0,0 +1,24 @@
+"""
+Tweek Agentic Memory
+
+Persistent, structured memory that enables Tweek to learn from past security
+decisions and make better screening choices over time.
+
+Features:
+- Pattern decision history with time-decay weighting
+- Source trustworthiness tracking (URL/file injection history)
+- Cross-session workflow baselines
+- Learned whitelist suggestions from approval patterns
+
+Safety Invariants:
+- CRITICAL+deterministic patterns are immune from memory adjustment
+- Memory can only relax ask -> log (never deny -> anything)
+- Project memory can escalate but never relax global decisions
+- Minimum 10 weighted decisions before any adjustment suggested
+- 30-day half-life for time decay
+- Full audit trail for every memory operation
+"""
+
+from tweek.memory.store import MemoryStore, get_memory_store
+
+__all__ = ["MemoryStore", "get_memory_store"]

tweek/memory/queries.py

@@ -0,0 +1,223 @@
+"""
+Tweek Memory Query Functions
+
+Hook entry points for reading and writing memory during PreToolUse
+and PostToolUse screening. All functions are best-effort and fail
+silently to avoid blocking security screening.
+"""
+
+from datetime import datetime
+from typing import Any, Dict, Optional
+
+from tweek.memory.safety import MIN_CONFIDENCE_SCORE
+from tweek.memory.schemas import PatternDecisionEntry
+from tweek.memory.store import (
+    MemoryStore,
+    content_hash,
+    get_memory_store,
+    hash_project,
+    normalize_path_prefix,
+)
+
+
+def memory_read_for_pattern(
+    pattern_name: str,
+    pattern_severity: str,
+    pattern_confidence: str,
+    tool_name: str,
+    path_prefix: Optional[str] = None,
+    project_hash: Optional[str] = None,
+    current_decision: str = "ask",
+) -> Optional[Dict[str, Any]]:
+    """Read memory for a pattern match to get confidence adjustment.
+
+    Called from PreToolUse after pattern matching, before enforcement resolution.
+
+    Returns a dict with 'adjusted_decision' key if memory suggests a change,
+    or None if no adjustment suggested.
+    """
+    try:
+        store = get_memory_store()
+        normalized_prefix = normalize_path_prefix(path_prefix) if path_prefix else None
+
+        adjustment = store.get_confidence_adjustment(
+            pattern_name=pattern_name,
+            path_prefix=normalized_prefix,
+            current_decision=current_decision,
+            original_severity=pattern_severity,
+            original_confidence=pattern_confidence,
+        )
+
+        if adjustment is None:
+            return None
+
+        if (
+            adjustment.adjusted_decision
+            and adjustment.confidence_score >= MIN_CONFIDENCE_SCORE
+        ):
+            return {
+                "adjusted_decision": adjustment.adjusted_decision,
+                "confidence_score": adjustment.confidence_score,
+                "approval_ratio": adjustment.approval_ratio,
+                "total_decisions": adjustment.total_decisions,
+                "pattern_name": pattern_name,
+            }
+
+        return None
+    except Exception:
+        return None
+
+
+def memory_write_after_decision(
+    pattern_name: str,
+    pattern_id: Optional[int],
+    original_severity: str,
+    original_confidence: str,
+    decision: str,
+    user_response: Optional[str],
+    tool_name: str,
+    content: str,
+    path_prefix: Optional[str] = None,
+    project_hash: Optional[str] = None,
+) -> None:
+    """Write a pattern decision to memory.
+
+    Called from PreToolUse after the decision is made (in all branches:
+    deny, ask, log, and allow).
+    """
+    try:
+        store = get_memory_store()
+        normalized_prefix = normalize_path_prefix(path_prefix) if path_prefix else None
+        c_hash = content_hash(content) if content else None
+
+        entry = PatternDecisionEntry(
+            pattern_name=pattern_name,
+            pattern_id=pattern_id,
+            original_severity=original_severity,
+            original_confidence=original_confidence,
+            decision=decision,
+            user_response=user_response,
+            tool_name=tool_name,
+            content_hash=c_hash,
+            path_prefix=normalized_prefix,
+            project_hash=project_hash,
+        )
+
+        store.record_decision(entry)
+    except Exception:
+        pass  # Memory is best-effort
+
+
+def memory_read_source_trust(
+    source_type: str,
+    source_key: str,
+) -> Optional[Dict[str, Any]]:
+    """Read source trust information.
+
+    Called from PostToolUse before screen_content().
+
+    Returns a dict with trust information if available.
+    """
+    try:
+        store = get_memory_store()
+        entry = store.get_source_trust(source_type, source_key)
+
+        if entry is None:
+            return None
+
+        return {
+            "source_type": entry.source_type,
+            "source_key": entry.source_key,
+            "trust_score": entry.trust_score,
+            "total_scans": entry.total_scans,
+            "injection_detections": entry.injection_detections,
+            "last_injection": entry.last_injection,
+        }
+    except Exception:
+        return None
+
+
+def memory_write_source_scan(
+    source_type: str,
+    source_key: str,
+    had_injection: bool,
+) -> None:
+    """Record a source scan result in memory.
+
+    Called from PostToolUse after screen_content() completes.
+    """
+    try:
+        store = get_memory_store()
+        store.record_source_scan(source_type, source_key, had_injection)
+
+        # Also record domain-level trust for URLs
+        if source_type == "url":
+            try:
+                from urllib.parse import urlparse
+                domain = urlparse(source_key).hostname
+                if domain:
+                    store.record_source_scan("domain", domain, had_injection)
+            except Exception:
+                pass
+    except Exception:
+        pass  # Memory is best-effort
+
+
+def memory_update_workflow(
+    project_hash: str,
+    tool_name: str,
+    was_denied: bool = False,
+) -> None:
+    """Update workflow baseline for a project.
+
+    Called from PreToolUse at the end of processing.
+    """
+    try:
+        store = get_memory_store()
+        hour = datetime.utcnow().hour
+        store.update_workflow(
+            project_hash=project_hash,
+            tool_name=tool_name,
+            hour_of_day=hour,
+            was_denied=was_denied,
+        )
+    except Exception:
+        pass  # Memory is best-effort
+
+
+def memory_get_workflow_baseline(
+    project_hash: str,
+) -> Optional[Dict[str, Any]]:
+    """Get workflow baseline for cross-session comparison.
+
+    Called from session_analyzer to compare current behavior against baselines.
+    """
+    try:
+        store = get_memory_store()
+        baselines = store.get_workflow_baseline(project_hash)
+
+        if not baselines:
+            return None
+
+        # Aggregate into a summary
+        tool_counts = {}
+        total_invocations = 0
+        total_denials = 0
+
+        for b in baselines:
+            if b.tool_name not in tool_counts:
+                tool_counts[b.tool_name] = {"invocations": 0, "denials": 0}
+            tool_counts[b.tool_name]["invocations"] += b.invocation_count
+            tool_counts[b.tool_name]["denials"] += b.denied_count
+            total_invocations += b.invocation_count
+            total_denials += b.denied_count
+
+        return {
+            "project_hash": project_hash,
+            "tool_counts": tool_counts,
+            "total_invocations": total_invocations,
+            "total_denials": total_denials,
+            "denial_ratio": total_denials / max(total_invocations, 1),
+        }
+    except Exception:
+        return None

tweek/memory/safety.py

@@ -0,0 +1,140 @@
+"""
+Tweek Memory Safety Module
+
+Enforces non-negotiable safety invariants for memory-based adjustments.
+
+Rules:
+1. CRITICAL+deterministic patterns are immune - memory NEVER adjusts them
+2. One-step max relaxation: ask -> log only (never deny -> anything)
+3. Additive-only project merge: project can escalate, never relax
+4. Minimum threshold: 10+ weighted decisions before any adjustment
+"""
+
+from typing import Optional
+
+# Patterns that are immune from any memory-based adjustment.
+# These are CRITICAL+deterministic patterns that should always deny.
+IMMUNE_SEVERITIES = frozenset({"critical"})
+IMMUNE_CONFIDENCES = frozenset({"deterministic"})
+
+# Decision hierarchy: deny > ask > log > allow
+DECISION_RANK = {"deny": 3, "ask": 2, "log": 1, "allow": 0}
+
+# Maximum relaxation: current_decision -> max allowed relaxation target
+# deny -> NOTHING (immune)
+# ask -> log (one step down)
+# log -> log (already minimum observable)
+MAX_RELAXATION = {
+    "deny": "deny",   # Never relax deny
+    "ask": "log",      # Can relax to log
+    "log": "log",      # Already at minimum
+    "allow": "allow",  # Already at minimum
+}
+
+# Minimum weighted decisions before memory can suggest adjustments
+MIN_DECISION_THRESHOLD = 10
+
+# Minimum approval ratio to suggest relaxation
+MIN_APPROVAL_RATIO = 0.90  # 90% approval rate
+
+# Minimum confidence score to actually apply an adjustment
+MIN_CONFIDENCE_SCORE = 0.80
+
+
+def is_immune_pattern(severity: str, confidence: str) -> bool:
+    """Check if a pattern is immune from memory adjustment.
+
+    CRITICAL+deterministic patterns are NEVER adjusted by memory.
+    This is a hard safety invariant enforced at every layer.
+    """
+    return (
+        severity.lower() in IMMUNE_SEVERITIES
+        and confidence.lower() in IMMUNE_CONFIDENCES
+    )
+
+
+def get_max_relaxation(current_decision: str) -> str:
+    """Get the maximum allowed relaxation target for a decision.
+
+    Returns the most relaxed decision memory is allowed to suggest.
+    """
+    return MAX_RELAXATION.get(current_decision, current_decision)
+
+
+def validate_memory_adjustment(
+    pattern_name: str,
+    original_severity: str,
+    original_confidence: str,
+    suggested_decision: str,
+    current_decision: str,
+) -> str:
+    """Validate and potentially apply a memory-suggested decision adjustment.
+
+    Returns the final decision after safety validation. This is the last
+    gate before a memory adjustment takes effect.
+
+    Args:
+        pattern_name: The pattern being evaluated
+        original_severity: Pattern's severity level
+        original_confidence: Pattern's confidence level
+        suggested_decision: What memory suggests
+        current_decision: The decision from enforcement policy
+
+    Returns:
+        The validated decision (may be same as current if adjustment rejected)
+    """
+    # Rule 1: CRITICAL+deterministic are immune
+    if is_immune_pattern(original_severity, original_confidence):
+        return current_decision
+
+    # Rule 2: deny is never relaxed by memory
+    if current_decision == "deny":
+        return current_decision
+
+    # Rule 3: Can only relax, never escalate via memory
+    # Memory is for reducing noise, not adding blocks
+    suggested_rank = DECISION_RANK.get(suggested_decision, 2)
+    current_rank = DECISION_RANK.get(current_decision, 2)
+    if suggested_rank >= current_rank:
+        # Suggested is same or stricter - no change needed
+        return current_decision
+
+    # Rule 4: Maximum one-step relaxation
+    max_relaxation = get_max_relaxation(current_decision)
+    max_rank = DECISION_RANK.get(max_relaxation, 2)
+    if suggested_rank < max_rank:
+        # Suggested goes beyond max relaxation - clamp to max
+        return max_relaxation
+
+    return suggested_decision
+
+
+def compute_suggested_decision(
+    current_decision: str,
+    approval_ratio: float,
+    total_weighted_decisions: float,
+    original_severity: str,
+    original_confidence: str,
+) -> Optional[str]:
+    """Compute what decision memory would suggest, if any.
+
+    Returns None if memory has no suggestion (insufficient data or
+    pattern is immune).
+    """
+    # Immune patterns get no suggestions
+    if is_immune_pattern(original_severity, original_confidence):
+        return None
+
+    # Insufficient data
+    if total_weighted_decisions < MIN_DECISION_THRESHOLD:
+        return None
+
+    # deny is never relaxed
+    if current_decision == "deny":
+        return None
+
+    # Only suggest relaxation if approval ratio is very high
+    if approval_ratio >= MIN_APPROVAL_RATIO and current_decision == "ask":
+        return "log"
+
+    return None

tweek/memory/schemas.py

@@ -0,0 +1,80 @@
+"""
+Tweek Memory Data Schemas
+
+Dataclasses for structured memory entries and query results.
+"""
+
+from dataclasses import dataclass, field
+from typing import Optional
+
+
+@dataclass
+class PatternDecisionEntry:
+    """A single pattern decision record."""
+
+    pattern_name: str
+    pattern_id: Optional[int]
+    original_severity: str
+    original_confidence: str
+    decision: str  # deny/ask/log/allow
+    user_response: Optional[str]  # approved/denied/null
+    tool_name: str
+    content_hash: Optional[str]
+    path_prefix: Optional[str]
+    project_hash: Optional[str]
+    timestamp: Optional[str] = None
+    decay_weight: float = 1.0
+
+
+@dataclass
+class ConfidenceAdjustment:
+    """Result of a memory confidence query for a pattern."""
+
+    pattern_name: str
+    path_prefix: Optional[str]
+    total_decisions: int
+    weighted_approvals: float
+    weighted_denials: float
+    approval_ratio: float
+    last_decision: Optional[str]
+    adjusted_decision: Optional[str] = None  # suggested decision override
+    confidence_score: float = 0.0  # 0.0-1.0 how confident the suggestion is
+
+
+@dataclass
+class SourceTrustEntry:
+    """Trust score for a URL, file, or domain."""
+
+    source_type: str  # url/file/domain
+    source_key: str
+    total_scans: int = 0
+    injection_detections: int = 0
+    trust_score: float = 0.5  # 0.0=bad, 1.0=good
+    last_clean_scan: Optional[str] = None
+    last_injection: Optional[str] = None
+
+
+@dataclass
+class WorkflowBaseline:
+    """Baseline tool usage pattern for a project."""
+
+    project_hash: str
+    tool_name: str
+    hour_of_day: Optional[int]
+    invocation_count: int = 0
+    denied_count: int = 0
+
+
+@dataclass
+class LearnedWhitelistSuggestion:
+    """A suggested whitelist entry derived from approval patterns."""
+
+    id: int
+    pattern_name: str
+    tool_name: Optional[str]
+    path_prefix: Optional[str]
+    approval_count: int = 0
+    denial_count: int = 0
+    confidence: float = 0.0
+    suggested_at: Optional[str] = None
+    human_reviewed: int = 0  # 0=pending, 1=accepted, -1=rejected