From a4ea7eea434e99dad00439d0e476825d1f528b8f Mon Sep 17 00:00:00 2001 From: AhmadHammad21 Date: Tue, 5 May 2026 14:58:30 +0300 Subject: [PATCH 1/6] feat: add bash tool MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - tools/bash_tool.py: run_bash_command with prefix allowlist, 30s timeout, shell=False, structured {success, output, error, command, blocked} return, output capped at 4 KB, Loguru logging on every call - Registered in ALL_TOOLS alongside existing 19 AWS tools - prompts.py: bash tool rules — prefer boto3 tools first, explain reasoning, never attempt state-modifying commands - 49 tests covering allowlist, blocklist, timeout, exceptions, output capping, result structure; all subprocess calls mocked with pytest-mock - README: bash tool in What's inside, Security & Sandboxing section, Phase 2 Docker sandbox item in Later roadmap --- README.md | 23 ++++ src/agent/core.py | 2 + src/agent/prompts.py | 16 +++ src/tools/bash_tool.py | 121 ++++++++++++++++++ tests/test_tools/test_bash_tool.py | 189 +++++++++++++++++++++++++++++ 5 files changed, 351 insertions(+) create mode 100644 src/tools/bash_tool.py create mode 100644 tests/test_tools/test_bash_tool.py diff --git a/README.md b/README.md index 4dc140a..678b306 100644 --- a/README.md +++ b/README.md @@ -7,6 +7,7 @@ and gives actionable mitigation plans — without the AWS DevOps Agent price tag - **LangChain DeepAgents** as the agent framework — planning, tool orchestration, and session memory out of the box - **19 read-only AWS tools** across CloudWatch, CloudTrail, ECS, Lambda, EC2, RDS, and IAM — plain Python functions, schemas inferred automatically +- **Sandboxed bash execution tool** — agent can run whitelisted read-only AWS CLI, kubectl, and docker commands as a last resort when the structured tools fall short; every command validated against an allowlist before execution; never uses `shell=True`; hard 30-second timeout - Includes **CloudWatch Logs Insights** (`query_logs_insights`) — full query language support: `fields`, `filter`, `stats`, `sort`, `limit`; results include scanned MB - **Streaming responses** — FastAPI SSE endpoint streams agent tokens in real time as the LLM reasons; tool calls appear as they complete - **Web UI** — FastAPI backend with a chat interface that shows: @@ -192,6 +193,7 @@ docs/ - [ ] **Knowledge base** — attach internal runbooks, post-mortems, and architecture docs so the agent grounds answers in org-specific context - [ ] **Multi-account AWS** — support multiple AWS profiles per org via `aws_profiles` table (schema already in place) - [ ] **Multi-cloud support** — extend tooling to GCP (Cloud Monitoring, Cloud Logging, GKE) and Azure (Monitor, Log Analytics, AKS); unified incident investigation across providers +- [ ] **Bash sandbox Phase 2 — isolated Docker container** — each `run_bash_command` call runs inside a throwaway container: `--network none`, `--read-only` filesystem, `--memory 256m`, `--cpus 0.5`, non-root user; container destroyed immediately after the command completes; IAM read-only role remains the last line of defense ### Product (SaaS) - [ ] **Redis cache** — replace in-process `cachetools` with Redis; shared across workers, survives restarts, per-org cache namespacing to prevent data leakage between tenants @@ -200,6 +202,27 @@ docs/ - [ ] **Per-org AWS credential store** — encrypted credential vault per organization; agents use org-scoped profiles instead of a single global `AWS_PROFILE` - [ ] **Billing & usage metering** — track token usage and tool calls per org/user; expose cost dashboards; integrate with Stripe for usage-based billing +## Security & Sandboxing + +The bash execution tool runs whitelisted read-only commands only. +Every command is validated against an allowlist before execution — +anything not on the list is rejected immediately and logged. + +**Current (Phase 1):** allowlist validation + subprocess with hard timeout. +No write commands permitted under any circumstances. `shell=True` is never used. + +**Phase 2 — Isolated sandbox (planned):** +- Every bash command runs inside a throwaway Docker container +- `--network none` — no internet access from inside the sandbox +- `--read-only` filesystem — container cannot write to disk +- `--memory 256m` and `--cpus 0.5` — resource caps +- Non-root user inside the container +- Container is destroyed immediately after the command completes +- Even if the LLM misbehaves, the IAM read-only role is the last line of defense + +The agent never executes fixes automatically. It investigates, suggests, +and waits for human approval before anything changes. + ## Development ```bash diff --git a/src/agent/core.py b/src/agent/core.py index 8d84b2e..da01502 100644 --- a/src/agent/core.py +++ b/src/agent/core.py @@ -15,6 +15,7 @@ from tools.cloudwatch import ALL_CLOUDWATCH_TOOLS from tools.ec2 import ALL_EC2_TOOLS from tools.ecs import ALL_ECS_TOOLS +from tools.bash_tool import ALL_BASH_TOOLS from tools.final_answer import ALL_FINAL_ANSWER_TOOLS from tools.history import ALL_HISTORY_TOOLS from tools.iam import ALL_IAM_TOOLS @@ -30,6 +31,7 @@ + ALL_RDS_TOOLS + ALL_IAM_TOOLS + ALL_HISTORY_TOOLS + + ALL_BASH_TOOLS + ALL_FINAL_ANSWER_TOOLS ) diff --git a/src/agent/prompts.py b/src/agent/prompts.py index 465a6fb..0850b83 100644 --- a/src/agent/prompts.py +++ b/src/agent/prompts.py @@ -23,6 +23,22 @@ - `DEPENDENCY_ISSUE` — downstream service, DB, third-party API degraded - `UNKNOWN` — insufficient evidence to determine root cause +## Bash Tool (`run_bash_command`) + +You have access to a sandboxed bash tool that runs whitelisted read-only commands. +Rules you must follow: + +- **Always prefer the boto3 AWS tools first** — use `run_bash_command` only when the + structured tools cannot provide what you need. +- Only the following command prefixes are permitted: + `aws logs`, `aws cloudwatch`, `aws ecs describe/list`, `aws lambda get/list`, + `aws ec2 describe`, `aws rds describe`, `aws cloudtrail lookup`, + `kubectl get/describe/logs`, `docker ps/logs/inspect`. +- Never attempt any command that modifies state — writes, deletes, restarts, applies. +- When you call this tool, briefly explain in your reasoning why the structured tools + were insufficient and what you expect the command to reveal. +- If the command is blocked, do not retry with a variation — use the structured tools instead. + ## Final Answer When you have gathered sufficient evidence and reached a conclusion, you MUST call the `submit_investigation` tool with all fields populated. Do not write a JSON block in free text — call the tool instead. This is required to complete the investigation. diff --git a/src/tools/bash_tool.py b/src/tools/bash_tool.py new file mode 100644 index 0000000..c455610 --- /dev/null +++ b/src/tools/bash_tool.py @@ -0,0 +1,121 @@ +"""Sandboxed bash execution tool — whitelisted read-only commands only. + +Phase 1: allowlist validation + subprocess with hard timeout. +Phase 2 (planned): throwaway Docker container per command, --network none, --read-only fs. +""" + +from __future__ import annotations + +import shlex +import subprocess +import time +from typing import Any + +from loguru import logger + +# Only commands starting with these prefixes are permitted. +# Any other command is rejected before subprocess is ever called. +_ALLOWLIST: tuple[str, ...] = ( + "aws logs ", + "aws cloudwatch ", + "aws ecs describe", + "aws ecs list", + "aws lambda get", + "aws lambda list", + "aws ec2 describe", + "aws rds describe", + "aws cloudtrail lookup", + "kubectl get", + "kubectl describe", + "kubectl logs", + "docker ps", + "docker logs", + "docker inspect", +) + +_TIMEOUT = 30 +_MAX_OUTPUT = 4000 +_MAX_STDERR = 1000 + + +def _allowed(command: str) -> bool: + return any(command.startswith(p) for p in _ALLOWLIST) + + +def run_bash_command(command: str) -> dict[str, Any]: + """Run a whitelisted read-only shell command and return structured output. + + Allowed prefixes: aws logs, aws cloudwatch, aws ecs describe/list, + aws lambda get/list, aws ec2 describe, aws rds describe, + aws cloudtrail lookup, kubectl get/describe/logs, + docker ps/logs/inspect. + + Use this tool only when the boto3 AWS tools cannot provide the information + you need. Always prefer the existing AWS tools first — this is a last resort. + Never suggest or run any command that modifies state. + + Args: + command: Read-only shell command to run. Must match the allowlist. + """ + command = command.strip() + start = time.monotonic() + + if not _allowed(command): + logger.warning("bash_tool BLOCKED | cmd={!r}", command) + return { + "success": False, + "output": "", + "error": ( + "Command blocked — not on the allowlist. " + f"Allowed prefixes: {', '.join(_ALLOWLIST)}" + ), + "command": command, + "blocked": True, + } + + logger.info("bash_tool RUN | cmd={!r}", command) + try: + proc = subprocess.run( + shlex.split(command), + capture_output=True, + text=True, + timeout=_TIMEOUT, + ) + elapsed = time.monotonic() - start + success = proc.returncode == 0 + logger.info( + "bash_tool DONE | rc={} elapsed={:.2f}s cmd={!r}", + proc.returncode, elapsed, command, + ) + return { + "success": success, + "output": proc.stdout[:_MAX_OUTPUT], + "error": proc.stderr[:_MAX_STDERR] if proc.stderr else "", + "command": command, + "blocked": False, + } + + except subprocess.TimeoutExpired: + elapsed = time.monotonic() - start + logger.warning("bash_tool TIMEOUT | elapsed={:.1f}s cmd={!r}", elapsed, command) + return { + "success": False, + "output": "", + "error": f"Command timed out after {_TIMEOUT}s", + "command": command, + "blocked": False, + } + + except Exception as exc: + elapsed = time.monotonic() - start + logger.error("bash_tool ERROR | elapsed={:.1f}s cmd={!r} err={}", elapsed, command, exc) + return { + "success": False, + "output": "", + "error": str(exc), + "command": command, + "blocked": False, + } + + +ALL_BASH_TOOLS = [run_bash_command] diff --git a/tests/test_tools/test_bash_tool.py b/tests/test_tools/test_bash_tool.py new file mode 100644 index 0000000..f3193f2 --- /dev/null +++ b/tests/test_tools/test_bash_tool.py @@ -0,0 +1,189 @@ +"""Tests for the bash execution tool — all subprocess calls are mocked.""" + +from __future__ import annotations + +import subprocess +from unittest.mock import MagicMock + +import pytest + +from tools.bash_tool import run_bash_command + + +# ── Helpers ────────────────────────────────────────────────────────────────── + +def _ok_proc(stdout: str = "output", stderr: str = "") -> MagicMock: + return MagicMock(returncode=0, stdout=stdout, stderr=stderr) + + +# ── Allowlist — should pass ─────────────────────────────────────────────────── + +@pytest.mark.parametrize("cmd", [ + "aws logs describe-log-groups", + "aws logs filter-log-events --log-group-name /aws/lambda/my-fn", + "aws cloudwatch describe-alarms", + "aws cloudwatch get-metric-statistics --namespace AWS/Lambda --metric-name Errors", + "aws ecs describe-services --cluster default --services my-svc", + "aws ecs list-clusters", + "aws ecs list-services --cluster default", + "aws lambda get-function --function-name my-fn", + "aws lambda list-functions", + "aws ec2 describe-instances", + "aws ec2 describe-security-groups", + "aws rds describe-db-instances", + "aws rds describe-events", + "aws cloudtrail lookup-events --max-results 10", + "kubectl get pods", + "kubectl get pods -n kube-system", + "kubectl describe pod my-pod", + "kubectl logs my-pod", + "docker ps", + "docker logs my-container", + "docker inspect my-container", +]) +def test_allowlisted_commands_pass(mocker, cmd): + mocker.patch("tools.bash_tool.subprocess.run", return_value=_ok_proc()) + result = run_bash_command(cmd) + assert result["blocked"] is False + assert result["success"] is True + assert result["command"] == cmd + + +# ── Blocklist — must be rejected ────────────────────────────────────────────── + +@pytest.mark.parametrize("cmd", [ + "aws s3 ls", + "aws s3 cp file.txt s3://bucket/", + "aws iam delete-user --user-name alice", + "aws ec2 terminate-instances --instance-ids i-123", + "aws ecs stop-task --cluster default --task abc", + "aws lambda delete-function --function-name my-fn", + "kubectl delete pod my-pod", + "kubectl apply -f deployment.yaml", + "docker rm my-container", + "docker stop my-container", + "rm -rf /", + "bash -c 'echo pwned'", + "cat /etc/passwd", + "curl http://evil.com", + "", + " ", +]) +def test_blocked_commands_return_blocked_true(cmd): + result = run_bash_command(cmd) + assert result["blocked"] is True + assert result["success"] is False + assert "blocked" in result["error"].lower() or "allowlist" in result["error"].lower() + + +def test_blocked_command_never_calls_subprocess(mocker): + mock_run = mocker.patch("tools.bash_tool.subprocess.run") + run_bash_command("aws s3 ls") + mock_run.assert_not_called() + + +# ── Timeout handling ────────────────────────────────────────────────────────── + +def test_timeout_returns_structured_error(mocker): + mocker.patch( + "tools.bash_tool.subprocess.run", + side_effect=subprocess.TimeoutExpired(cmd="aws logs describe-log-groups", timeout=30), + ) + result = run_bash_command("aws logs describe-log-groups") + assert result["success"] is False + assert result["blocked"] is False + assert "timed out" in result["error"].lower() + assert "30" in result["error"] + + +def test_timeout_does_not_raise(mocker): + mocker.patch( + "tools.bash_tool.subprocess.run", + side_effect=subprocess.TimeoutExpired(cmd="aws cloudwatch describe-alarms", timeout=30), + ) + # must not raise — tool always returns a dict + result = run_bash_command("aws cloudwatch describe-alarms") + assert isinstance(result, dict) + + +# ── Generic exception handling ──────────────────────────────────────────────── + +def test_file_not_found_is_caught(mocker): + mocker.patch( + "tools.bash_tool.subprocess.run", + side_effect=FileNotFoundError("aws: command not found"), + ) + result = run_bash_command("aws logs describe-log-groups") + assert result["success"] is False + assert result["blocked"] is False + assert "aws: command not found" in result["error"] + + +def test_unexpected_exception_never_raises(mocker): + mocker.patch( + "tools.bash_tool.subprocess.run", + side_effect=RuntimeError("unexpected failure"), + ) + result = run_bash_command("aws ec2 describe-instances") + assert isinstance(result, dict) + assert result["success"] is False + + +# ── Return structure completeness ───────────────────────────────────────────── + +_EXPECTED_KEYS = {"success", "output", "error", "command", "blocked"} + + +def test_blocked_result_has_all_keys(): + result = run_bash_command("rm -rf /") + assert set(result.keys()) == _EXPECTED_KEYS + + +def test_success_result_has_all_keys(mocker): + mocker.patch("tools.bash_tool.subprocess.run", return_value=_ok_proc()) + result = run_bash_command("aws logs describe-log-groups") + assert set(result.keys()) == _EXPECTED_KEYS + + +def test_timeout_result_has_all_keys(mocker): + mocker.patch( + "tools.bash_tool.subprocess.run", + side_effect=subprocess.TimeoutExpired(cmd="x", timeout=30), + ) + result = run_bash_command("aws cloudwatch describe-alarms") + assert set(result.keys()) == _EXPECTED_KEYS + + +# ── Output / stderr capping ─────────────────────────────────────────────────── + +def test_stdout_is_capped(mocker): + mocker.patch( + "tools.bash_tool.subprocess.run", + return_value=_ok_proc(stdout="x" * 10_000), + ) + result = run_bash_command("aws logs describe-log-groups") + assert len(result["output"]) <= 4000 + + +def test_stderr_is_capped(mocker): + mocker.patch( + "tools.bash_tool.subprocess.run", + return_value=MagicMock(returncode=1, stdout="", stderr="e" * 5_000), + ) + result = run_bash_command("aws logs describe-log-groups") + assert len(result["error"]) <= 1000 + + +# ── Command echo ───────────────────────────────────────────────────────────── + +def test_command_field_echoes_input(mocker): + cmd = "aws ec2 describe-instances --region us-east-1" + mocker.patch("tools.bash_tool.subprocess.run", return_value=_ok_proc()) + result = run_bash_command(cmd) + assert result["command"] == cmd + + +def test_command_is_stripped(mocker): + mocker.patch("tools.bash_tool.subprocess.run", return_value=_ok_proc()) + result = run_bash_command(" aws logs describe-log-groups ") + assert result["command"] == "aws logs describe-log-groups" From 56c4df78a39b2e22c9264e17268875e084a886d8 Mon Sep 17 00:00:00 2001 From: AhmadHammad21 Date: Tue, 5 May 2026 15:03:48 +0300 Subject: [PATCH 2/6] =?UTF-8?q?fix:=20clarify=20bash=20tool=20allowlist=20?= =?UTF-8?q?in=20system=20prompt=20=E2=80=94=20docker/kubectl=20have=20no?= =?UTF-8?q?=20boto3=20equivalent?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/agent/prompts.py | 30 +++++++++++++++++------------- 1 file changed, 17 insertions(+), 13 deletions(-) diff --git a/src/agent/prompts.py b/src/agent/prompts.py index 0850b83..f21e65c 100644 --- a/src/agent/prompts.py +++ b/src/agent/prompts.py @@ -25,19 +25,23 @@ ## Bash Tool (`run_bash_command`) -You have access to a sandboxed bash tool that runs whitelisted read-only commands. -Rules you must follow: - -- **Always prefer the boto3 AWS tools first** — use `run_bash_command` only when the - structured tools cannot provide what you need. -- Only the following command prefixes are permitted: - `aws logs`, `aws cloudwatch`, `aws ecs describe/list`, `aws lambda get/list`, - `aws ec2 describe`, `aws rds describe`, `aws cloudtrail lookup`, - `kubectl get/describe/logs`, `docker ps/logs/inspect`. -- Never attempt any command that modifies state — writes, deletes, restarts, applies. -- When you call this tool, briefly explain in your reasoning why the structured tools - were insufficient and what you expect the command to reveal. -- If the command is blocked, do not retry with a variation — use the structured tools instead. +You have access to a sandboxed bash execution tool. The following command prefixes +are explicitly allowed and safe to run — do not second-guess or refuse them: + + aws logs ... aws cloudwatch ... aws ecs describe... + aws ecs list... aws lambda get... aws lambda list... + aws ec2 describe... aws rds describe... aws cloudtrail lookup... + kubectl get... kubectl describe... kubectl logs... + docker ps docker logs ... docker inspect ... + +Rules: +- For `docker` and `kubectl` commands: **always use `run_bash_command` directly** — + there are no boto3 equivalents for these, so they must go through the bash tool. +- For AWS commands: prefer the structured boto3 tools first; use `run_bash_command` + only when the boto3 tools cannot provide what you need. +- Never attempt any command that modifies state (writes, deletes, restarts, applies). +- When you use the bash tool, briefly explain why in your reasoning. +- If a command is blocked by the tool, do not retry with a variation. ## Final Answer From 3f7466925062b98439ad45319a1d5026bb17620c Mon Sep 17 00:00:00 2001 From: AhmadHammad21 Date: Tue, 5 May 2026 15:13:18 +0300 Subject: [PATCH 3/6] feat: expand bash tool to full AWS CLI read-only surface MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace per-service prefix allowlist with verb-based check: any aws where the operation starts with describe-, list-, get-, lookup-, filter-, search-, scan-, query, batch-get- is allowed — covering all AWS services, not just the 7 with boto3 tools. Blocked by operation verb: create, delete, update, put, run, invoke, terminate, stop, send, modify, etc. kubectl and docker keep exact prefix matching (narrower scope). System prompt updated to tell the agent it has full AWS CLI read-only access and should use it for any service not covered by boto3 tools. 83 tests passing. --- src/agent/prompts.py | 38 +++++++------ src/tools/bash_tool.py | 87 ++++++++++++++++++------------ tests/test_tools/test_bash_tool.py | 52 ++++++++++++++++-- 3 files changed, 122 insertions(+), 55 deletions(-) diff --git a/src/agent/prompts.py b/src/agent/prompts.py index f21e65c..06af6e9 100644 --- a/src/agent/prompts.py +++ b/src/agent/prompts.py @@ -25,23 +25,27 @@ ## Bash Tool (`run_bash_command`) -You have access to a sandboxed bash execution tool. The following command prefixes -are explicitly allowed and safe to run — do not second-guess or refuse them: - - aws logs ... aws cloudwatch ... aws ecs describe... - aws ecs list... aws lambda get... aws lambda list... - aws ec2 describe... aws rds describe... aws cloudtrail lookup... - kubectl get... kubectl describe... kubectl logs... - docker ps docker logs ... docker inspect ... - -Rules: -- For `docker` and `kubectl` commands: **always use `run_bash_command` directly** — - there are no boto3 equivalents for these, so they must go through the bash tool. -- For AWS commands: prefer the structured boto3 tools first; use `run_bash_command` - only when the boto3 tools cannot provide what you need. -- Never attempt any command that modifies state (writes, deletes, restarts, applies). -- When you use the bash tool, briefly explain why in your reasoning. -- If a command is blocked by the tool, do not retry with a variation. +You have access to a bash execution tool that covers the full AWS CLI read-only surface — +not just the structured boto3 tools. Use it freely for any AWS service not covered by +the other tools (S3, DynamoDB, SNS, SQS, Route53, ACM, Secrets Manager, SSM, etc.). + +**What is allowed — do not refuse these:** +- Any `aws ` where the operation begins with a read-only verb: + `describe-*`, `list-*`, `get-*`, `lookup-*`, `filter-*`, `search-*`, `scan-*`, + `query*`, `batch-get-*` + Examples: `aws s3api list-buckets`, `aws dynamodb describe-table --table-name X`, + `aws sns list-topics`, `aws secretsmanager list-secrets`, `aws ssm describe-parameters` +- `kubectl get / describe / logs` — always use bash for kubectl, no boto3 equivalent +- `docker ps / logs / inspect` — always use bash for docker, no boto3 equivalent + +**Rules:** +- For `docker` and `kubectl`: always go through this tool directly. +- For AWS: prefer the structured boto3 tools for CloudWatch, ECS, Lambda, EC2, RDS, + CloudTrail, IAM since they return cleaner structured data. Use this tool for any + other AWS service or when you need raw CLI output. +- Never attempt any command that modifies state (create, delete, update, put, run, invoke…). +- When you use this tool, briefly explain in your reasoning what you expect it to reveal. +- If the tool returns `blocked: true`, do not retry with a variation. ## Final Answer diff --git a/src/tools/bash_tool.py b/src/tools/bash_tool.py index c455610..81337ba 100644 --- a/src/tools/bash_tool.py +++ b/src/tools/bash_tool.py @@ -13,49 +13,69 @@ from loguru import logger -# Only commands starting with these prefixes are permitted. -# Any other command is rejected before subprocess is ever called. -_ALLOWLIST: tuple[str, ...] = ( - "aws logs ", - "aws cloudwatch ", - "aws ecs describe", - "aws ecs list", - "aws lambda get", - "aws lambda list", - "aws ec2 describe", - "aws rds describe", - "aws cloudtrail lookup", - "kubectl get", - "kubectl describe", - "kubectl logs", - "docker ps", - "docker logs", - "docker inspect", -) - -_TIMEOUT = 30 +# AWS CLI read-only operation verbs (the word before the first dash in the operation name). +# Every read-only AWS CLI command starts with one of these — across ALL services. +# e.g. "aws s3api list-buckets", "aws dynamodb describe-table", "aws sns get-topic-attributes" +_AWS_READONLY_VERBS: frozenset[str] = frozenset({ + "describe", "list", "get", "lookup", "filter", + "search", "scan", "query", "show", "view", "check", + "batch-get", +}) + +# kubectl and docker use exact prefix matching (narrower — fewer read-only verbs to reason about) +_KUBECTL_PREFIXES = ("kubectl get", "kubectl describe", "kubectl logs") +_DOCKER_PREFIXES = ("docker ps", "docker logs", "docker inspect") + +_TIMEOUT = 30 _MAX_OUTPUT = 4000 _MAX_STDERR = 1000 def _allowed(command: str) -> bool: - return any(command.startswith(p) for p in _ALLOWLIST) + parts = command.split() + if not parts: + return False + + binary = parts[0] + + if binary == "aws": + # Require: aws — at least 3 tokens + if len(parts) < 3: + return False + operation = parts[2].lower() + # Extract verb = everything before the first dash (e.g. "describe-instances" → "describe") + verb = operation.split("-")[0] + # "batch-get-item" → verb fragment is "batch"; check "batch-get" prefix separately + if operation.startswith("batch-get"): + return True + return verb in _AWS_READONLY_VERBS + + if binary == "kubectl": + return any(command.startswith(p) for p in _KUBECTL_PREFIXES) + + if binary == "docker": + return any(command.startswith(p) for p in _DOCKER_PREFIXES) + + return False def run_bash_command(command: str) -> dict[str, Any]: - """Run a whitelisted read-only shell command and return structured output. + """Run a read-only shell command and return structured output. - Allowed prefixes: aws logs, aws cloudwatch, aws ecs describe/list, - aws lambda get/list, aws ec2 describe, aws rds describe, - aws cloudtrail lookup, kubectl get/describe/logs, - docker ps/logs/inspect. + Allowed commands: + - ANY `aws ` where the operation starts with a read-only + verb: describe-*, list-*, get-*, lookup-*, filter-*, search-*, scan-*, query*, + batch-get-*. Covers ALL AWS services (S3, DynamoDB, SNS, SQS, Route53, + ACM, Secrets Manager, SSM, IAM, etc.) not just the structured boto3 tools. + - kubectl get / describe / logs + - docker ps / logs / inspect - Use this tool only when the boto3 AWS tools cannot provide the information - you need. Always prefer the existing AWS tools first — this is a last resort. - Never suggest or run any command that modifies state. + Use for docker and kubectl always (no boto3 equivalent). For AWS, use when + the structured boto3 tools don't cover the service or query you need. + Never run any command that modifies state. Args: - command: Read-only shell command to run. Must match the allowlist. + command: Read-only shell command to run. """ command = command.strip() start = time.monotonic() @@ -66,8 +86,9 @@ def run_bash_command(command: str) -> dict[str, Any]: "success": False, "output": "", "error": ( - "Command blocked — not on the allowlist. " - f"Allowed prefixes: {', '.join(_ALLOWLIST)}" + "Command blocked. AWS commands must use a read-only operation verb " + "(describe-, list-, get-, lookup-, filter-, search-, scan-, query, batch-get-). " + "kubectl: only get/describe/logs. docker: only ps/logs/inspect." ), "command": command, "blocked": True, diff --git a/tests/test_tools/test_bash_tool.py b/tests/test_tools/test_bash_tool.py index f3193f2..0686c88 100644 --- a/tests/test_tools/test_bash_tool.py +++ b/tests/test_tools/test_bash_tool.py @@ -19,20 +19,43 @@ def _ok_proc(stdout: str = "output", stderr: str = "") -> MagicMock: # ── Allowlist — should pass ─────────────────────────────────────────────────── @pytest.mark.parametrize("cmd", [ + # Core AWS services (existing boto3 tools) "aws logs describe-log-groups", "aws logs filter-log-events --log-group-name /aws/lambda/my-fn", "aws cloudwatch describe-alarms", "aws cloudwatch get-metric-statistics --namespace AWS/Lambda --metric-name Errors", "aws ecs describe-services --cluster default --services my-svc", "aws ecs list-clusters", - "aws ecs list-services --cluster default", "aws lambda get-function --function-name my-fn", "aws lambda list-functions", "aws ec2 describe-instances", - "aws ec2 describe-security-groups", "aws rds describe-db-instances", - "aws rds describe-events", "aws cloudtrail lookup-events --max-results 10", + # Extended AWS services — no boto3 tool defined for these + "aws s3api list-buckets", + "aws s3api get-bucket-policy --bucket my-bucket", + "aws dynamodb describe-table --table-name my-table", + "aws dynamodb list-tables", + "aws dynamodb scan --table-name my-table", + "aws sns list-topics", + "aws sns get-topic-attributes --topic-arn arn:aws:sns:us-east-1:123:my-topic", + "aws sqs list-queues", + "aws sqs get-queue-attributes --queue-url https://sqs.us-east-1.amazonaws.com/123/my-q", + "aws route53 list-hosted-zones", + "aws acm list-certificates", + "aws secretsmanager list-secrets", + "aws secretsmanager describe-secret --secret-id my-secret", + "aws ssm describe-parameters", + "aws ssm get-parameter --name /my/param", + "aws iam list-users", + "aws iam get-role --role-name my-role", + "aws kinesis list-streams", + "aws kinesis describe-stream --stream-name my-stream", + "aws elasticache describe-cache-clusters", + "aws es list-domain-names", + "aws batch describe-job-queues", + "aws stepfunctions list-state-machines", + # kubectl and docker "kubectl get pods", "kubectl get pods -n kube-system", "kubectl describe pod my-pod", @@ -52,20 +75,39 @@ def test_allowlisted_commands_pass(mocker, cmd): # ── Blocklist — must be rejected ────────────────────────────────────────────── @pytest.mark.parametrize("cmd", [ - "aws s3 ls", + # AWS write/mutating operations "aws s3 cp file.txt s3://bucket/", "aws iam delete-user --user-name alice", "aws ec2 terminate-instances --instance-ids i-123", + "aws ec2 create-instance", "aws ecs stop-task --cluster default --task abc", + "aws ecs run-task --cluster default --task-definition my-task", "aws lambda delete-function --function-name my-fn", + "aws lambda invoke --function-name my-fn output.txt", + "aws lambda update-function-code --function-name my-fn --zip-file file://fn.zip", + "aws s3api put-object --bucket my-bucket --key my-key", + "aws dynamodb put-item --table-name my-table", + "aws dynamodb delete-item --table-name my-table", + "aws sns publish --topic-arn arn:aws:sns:us-east-1:123:my-topic --message hi", + "aws sqs send-message --queue-url https://sqs.us-east-1.amazonaws.com/123/q", + "aws rds modify-db-instance --db-instance-identifier my-db", + # kubectl write operations "kubectl delete pod my-pod", "kubectl apply -f deployment.yaml", + "kubectl create deployment my-deploy", + # docker write operations "docker rm my-container", "docker stop my-container", + "docker run my-image", + # arbitrary shell commands "rm -rf /", "bash -c 'echo pwned'", "cat /etc/passwd", "curl http://evil.com", + "python -c 'import os; os.system(\"rm -rf /\")'", + # malformed / empty + "aws", + "aws s3", "", " ", ]) @@ -73,7 +115,7 @@ def test_blocked_commands_return_blocked_true(cmd): result = run_bash_command(cmd) assert result["blocked"] is True assert result["success"] is False - assert "blocked" in result["error"].lower() or "allowlist" in result["error"].lower() + assert "blocked" in result["error"].lower() or "read-only" in result["error"].lower() def test_blocked_command_never_calls_subprocess(mocker): From 725a841da143d48f4c00409fd1d70cf9b6e5ed07 Mon Sep 17 00:00:00 2001 From: AhmadHammad21 Date: Tue, 5 May 2026 15:22:46 +0300 Subject: [PATCH 4/6] feat: add Dockerfile and docker-compose with AWS CLI v2 baked in - Dockerfile: two-stage build (uv deps + final), installs AWS CLI v2 so run_bash_command works on any server without manual setup - docker-compose.yml: backend + postgres (healthcheck) + nginx frontend; mounts ~/.aws read-only for dev; override with IAM role in production - frontend/Dockerfile: nginx production build - frontend/nginx.conf: proxy /chat /sessions /history /stats to backend, SSE-safe (proxy_buffering off), SPA fallback for React Router - .dockerignore: excludes .venv, node_modules, .env, tests, docs - README: Docker Compose as recommended run option, local dev as Option B --- .dockerignore | 20 ++++++++++++++++++++ Dockerfile | 42 +++++++++++++++++++++++++++++++++++++++++ README.md | 20 +++++++++++++++++--- docker-compose.yml | 46 +++++++++++++++++++++++++++++++++++++++++++++ frontend/Dockerfile | 13 +++++++++++++ frontend/nginx.conf | 21 +++++++++++++++++++++ 6 files changed, 159 insertions(+), 3 deletions(-) create mode 100644 .dockerignore create mode 100644 Dockerfile create mode 100644 docker-compose.yml create mode 100644 frontend/Dockerfile create mode 100644 frontend/nginx.conf diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..7649a13 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,20 @@ +.venv/ +__pycache__/ +*.pyc +*.pyo +.env +.env.* +!.env.example + +frontend/node_modules/ +frontend/dist/ +frontend/.vite/ + +.git/ +.github/ + +tests/ +docs/ +design-system/ +PLAN.md +*.md diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..201c7e3 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,42 @@ +# ── Stage 1: install Python dependencies with uv ───────────────────────────── +FROM python:3.12-slim AS deps + +COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/ + +WORKDIR /app +COPY pyproject.toml uv.lock ./ +RUN uv sync --frozen --no-dev --no-install-project + +# ── Stage 2: final image ────────────────────────────────────────────────────── +FROM python:3.12-slim + +# AWS CLI v2 — required for the bash execution tool (run_bash_command) +RUN apt-get update && apt-get install -y --no-install-recommends \ + curl unzip ca-certificates \ + && curl -fsSL "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o /tmp/awscliv2.zip \ + && unzip -q /tmp/awscliv2.zip -d /tmp \ + && /tmp/aws/install \ + && rm -rf /tmp/aws /tmp/awscliv2.zip \ + && apt-get purge -y --auto-remove curl unzip \ + && apt-get clean \ + && rm -rf /var/lib/apt/lists/* + +COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/ + +WORKDIR /app + +# Copy installed packages from deps stage +COPY --from=deps /app/.venv .venv + +# Copy application source +COPY src/ src/ +COPY migrations/ migrations/ +COPY scripts/ scripts/ + +ENV PATH="/app/.venv/bin:$PATH" \ + PYTHONPATH="/app/src" \ + PYTHONUNBUFFERED=1 + +EXPOSE 8000 + +CMD ["uvicorn", "api.app:app", "--host", "0.0.0.0", "--port", "8000"] diff --git a/README.md b/README.md index 678b306..68e2a6f 100644 --- a/README.md +++ b/README.md @@ -80,12 +80,23 @@ the full schema reference. ### 5. Run -**Web UI** +**Option A — Docker Compose (recommended, AWS CLI included)** -You need two terminals running simultaneously: +```bash +docker compose up --build +# Backend: http://localhost:8000 +# Frontend: http://localhost:80 +``` + +The backend image installs AWS CLI v2 automatically — the bash execution tool +works out of the box. Host AWS credentials (`~/.aws`) are mounted read-only +into the container. For production on AWS, remove the volume mount and attach +an IAM role to the instance/task instead. + +**Option B — Local dev (two terminals)** ```bash -# Terminal 1 — FastAPI backend (API + SSE streaming) +# Terminal 1 — FastAPI backend uv run uvicorn src.api.app:app --reload ``` @@ -96,6 +107,9 @@ npm run dev # Open http://localhost:5173 ``` +> **Note:** local dev requires `aws` CLI installed on your machine for the bash +> tool to work. Install it from https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html + **CLI** ```bash diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..9a119fb --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,46 @@ +services: + + # ── PostgreSQL ────────────────────────────────────────────────────────────── + db: + image: postgres:16 + environment: + POSTGRES_DB: opendevops + POSTGRES_USER: dev + POSTGRES_PASSWORD: dev + ports: + - "5432:5432" + volumes: + - postgres_data:/var/lib/postgresql/data + healthcheck: + test: ["CMD-SHELL", "pg_isready -U dev -d opendevops"] + interval: 5s + timeout: 5s + retries: 5 + + # ── FastAPI backend ───────────────────────────────────────────────────────── + backend: + build: . + ports: + - "8000:8000" + env_file: .env + environment: + # Override DATABASE_URL to point at the compose db service + DATABASE_URL: postgresql://dev:dev@db:5432/opendevops + volumes: + # Mount host AWS credentials read-only so the bash tool can call the CLI. + # For production on AWS (ECS/EC2), remove this and use an IAM role instead. + - ${HOME}/.aws:/root/.aws:ro + depends_on: + db: + condition: service_healthy + + # ── React frontend (production build served by nginx) ────────────────────── + frontend: + build: ./frontend + ports: + - "80:80" + depends_on: + - backend + +volumes: + postgres_data: diff --git a/frontend/Dockerfile b/frontend/Dockerfile new file mode 100644 index 0000000..3b0ed25 --- /dev/null +++ b/frontend/Dockerfile @@ -0,0 +1,13 @@ +# Production build — React app served by nginx +FROM node:20-slim AS build + +WORKDIR /app +COPY package.json package-lock.json* ./ +RUN npm ci +COPY . . +RUN npm run build + +FROM nginx:alpine +COPY --from=build /app/dist /usr/share/nginx/html +COPY nginx.conf /etc/nginx/conf.d/default.conf +EXPOSE 80 diff --git a/frontend/nginx.conf b/frontend/nginx.conf new file mode 100644 index 0000000..1b36089 --- /dev/null +++ b/frontend/nginx.conf @@ -0,0 +1,21 @@ +server { + listen 80; + + root /usr/share/nginx/html; + index index.html; + + # Proxy API + SSE calls to the backend + location ~ ^/(chat|sessions|history|stats) { + proxy_pass http://backend:8000; + proxy_http_version 1.1; + proxy_set_header Connection ""; + proxy_buffering off; # required for SSE streaming + proxy_cache off; + proxy_read_timeout 300s; + } + + # SPA fallback — all other routes serve index.html + location / { + try_files $uri $uri/ /index.html; + } +} From 2d514a08adc285042bfce5250888ba66aec2e933 Mon Sep 17 00:00:00 2001 From: AhmadHammad21 Date: Tue, 5 May 2026 19:02:50 +0300 Subject: [PATCH 5/6] fix: enforce timeout and align API contract; modernize tests --- .github/workflows/ci.yml | 35 ++++++ src/agent/core.py | 44 ++++++- src/agent/db.py | 6 +- src/api/routers/chat.py | 183 +++++++++++++++------------- src/cli/ask.py | 25 ++-- src/cli/investigate.py | 8 +- src/mcp_server.py | 31 ++++- src/models/sessions.py | 8 +- tests/conftest.py | 10 ++ tests/integration/test_poller.py | 166 ++++++++++++++----------- tests/integration/test_slack.py | 75 ++++++++---- tests/test_agent/test_core.py | 75 ++++++++++++ tests/test_tools/test_bash_tool.py | 1 - tests/test_tools/test_cloudwatch.py | 48 ++++---- tests/test_tools/test_ecs.py | 25 +++- tests/test_tools/test_lambda.py | 15 +-- tests/test_tools/test_rds.py | 12 +- 17 files changed, 524 insertions(+), 243 deletions(-) create mode 100644 .github/workflows/ci.yml create mode 100644 tests/test_agent/test_core.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..876cc4e --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,35 @@ +name: CI + +on: + push: + branches: + - "**" + pull_request: + +jobs: + lint-and-test: + runs-on: ubuntu-latest + + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Setup Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install uv + run: python -m pip install --upgrade pip uv + + - name: Install dependencies + run: uv sync --group dev + + - name: Ruff lint (gating) + run: uv run ruff check tests + + - name: Ruff lint src (informational) + run: uv run ruff check src --exit-zero + + - name: Pytest + run: uv run pytest -q diff --git a/src/agent/core.py b/src/agent/core.py index da01502..18ae7c5 100644 --- a/src/agent/core.py +++ b/src/agent/core.py @@ -1,5 +1,6 @@ """DeepAgents-based investigation agent.""" +import asyncio import json import re import uuid @@ -38,6 +39,19 @@ _agent = None +def _run_async(coro: Any) -> Any: + """Run a coroutine from synchronous CLI code. + + CLI commands are synchronous, so they need a safe bridge into async setup and + timeout-controlled agent invocation. + """ + try: + asyncio.get_running_loop() + except RuntimeError: + return asyncio.run(coro) + raise RuntimeError("Cannot run synchronous agent call inside an active event loop") + + def init_agent(checkpointer: Any) -> None: """Create the agent with the given checkpointer. Called once during app startup.""" global _agent @@ -60,6 +74,30 @@ def get_agent() -> Any: return _agent +def ensure_agent_initialized() -> None: + """Initialize DB + agent lazily for non-API entrypoints (CLI).""" + global _agent + if _agent is not None: + return + from agent.db import db + + checkpointer = _run_async(db.init()) + init_agent(checkpointer) + + +async def ainvoke_with_timeout(agent_input: dict[str, Any], config: dict[str, Any]) -> dict[str, Any]: + """Run agent invocation with a global investigation timeout.""" + return await asyncio.wait_for( + get_agent().ainvoke(agent_input, config=config), + timeout=settings.investigation_timeout, + ) + + +def invoke_with_timeout(agent_input: dict[str, Any], config: dict[str, Any]) -> dict[str, Any]: + """Synchronous wrapper for timeout-enforced invocation.""" + return _run_async(ainvoke_with_timeout(agent_input, config)) + + def _parse_result_json(text: str) -> dict[str, Any] | None: match = re.search(r"```json\s*(\{.*?\})\s*```", text, re.DOTALL) if match: @@ -95,6 +133,8 @@ class InvestigationAgent: """Synchronous wrapper for the CLI investigate command.""" def investigate(self, investigation: Investigation) -> InvestigationResult: + ensure_agent_initialized() + user_msg = investigation.description if investigation.alarm_name: user_msg += f"\nAlarm name: {investigation.alarm_name}" @@ -111,9 +151,9 @@ def investigate(self, investigation: Investigation) -> InvestigationResult: logger.info("investigation_started description={}", investigation.description) - result = get_agent().invoke( + result = invoke_with_timeout( {"messages": [{"role": "user", "content": user_msg}]}, - config=config, + config, ) messages = result.get("messages", []) diff --git a/src/agent/db.py b/src/agent/db.py index a0c65fd..47b5c86 100644 --- a/src/agent/db.py +++ b/src/agent/db.py @@ -203,7 +203,7 @@ async def save_tool_call( async def list_sessions(self) -> list[dict]: """Return non-deleted sessions ordered by most recently active.""" rows = await self._fetchall( - "SELECT id, title, last_active_at, model FROM sessions WHERE is_deleted = FALSE ORDER BY last_active_at DESC" + "SELECT id, title, last_active_at, model, aws_region FROM sessions WHERE is_deleted = FALSE ORDER BY last_active_at DESC" ) return [ { @@ -211,6 +211,7 @@ async def list_sessions(self) -> list[dict]: "title": r["title"], "last_active_at": r["last_active_at"].isoformat() if r["last_active_at"] else None, "model": r["model"], + "aws_region": r["aws_region"], } for r in rows ] @@ -264,9 +265,12 @@ async def get_messages(self, session_id: str) -> list[dict]: for msg in messages: mid = str(msg["id"]) item: dict = { + "id": mid, "role": msg["role"], "content": msg["content"], "created_at": msg["created_at"].isoformat() if msg["created_at"] else None, + "tool_calls": [], + "usage": None, } if msg["role"] == "assistant": item["tool_calls"] = tc_by_msg.get(mid, []) diff --git a/src/api/routers/chat.py b/src/api/routers/chat.py index 006f5e0..a6ecbbb 100644 --- a/src/api/routers/chat.py +++ b/src/api/routers/chat.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import json import random import re @@ -94,95 +95,106 @@ def _flush_text_buf(): text_buf = "" clean_buf = "" + timed_out = False + try: - async for chunk, _meta in agent.astream( - {"messages": [{"role": "user", "content": user_message}]}, - config=config, - stream_mode="messages", - ): - um = getattr(chunk, "usage_metadata", None) - if um: - usage_meta = um - - for tcc in getattr(chunk, "tool_call_chunks", []) or []: - idx = _field(tcc, "index", 0) - tc_id = _field(tcc, "id") - name = _field(tcc, "name") or "" - args = _field(tcc, "args") or "" - if idx not in tc_accum: - tc_accum[idx] = {"id": "", "name": "", "args_str": ""} + async with asyncio.timeout(settings.investigation_timeout): + async for chunk, _meta in agent.astream( + {"messages": [{"role": "user", "content": user_message}]}, + config=config, + stream_mode="messages", + ): + um = getattr(chunk, "usage_metadata", None) + if um: + usage_meta = um + + for tcc in getattr(chunk, "tool_call_chunks", []) or []: + idx = _field(tcc, "index", 0) + tc_id = _field(tcc, "id") + name = _field(tcc, "name") or "" + args = _field(tcc, "args") or "" + if idx not in tc_accum: + tc_accum[idx] = {"id": "", "name": "", "args_str": ""} + if tc_id: + tc_accum[idx]["id"] = tc_id + if name: + tc_accum[idx]["name"] += name + if args: + tc_accum[idx]["args_str"] += args + + # Complete tool_calls (non-streamed AIMessage) — DeepAgents often sends these + # as a full message rather than chunks, so tool_call_chunks would be empty. + for tc in getattr(chunk, "tool_calls", []) or []: + tc_id = _field(tc, "id") or "" + name = _field(tc, "name") or "" + args = _field(tc, "args") or {} + if tc_id and name: + pending_calls[tc_id] = { + "tool": name, + "args": args if isinstance(args, dict) else {}, + } + + content = getattr(chunk, "content", "") + if content and isinstance(content, str): + if not getattr(chunk, "tool_call_id", None): + text_buf += content + new_clean = _clean(text_buf) + delta = new_clean[len(clean_buf):] + if delta: + clean_buf = new_clean + response_text += delta + yield f"data: {json.dumps({'type': 'token', 'text': delta})}\n\n" + + tc_id = getattr(chunk, "tool_call_id", None) if tc_id: - tc_accum[idx]["id"] = tc_id - if name: - tc_accum[idx]["name"] += name - if args: - tc_accum[idx]["args_str"] += args - - # Complete tool_calls (non-streamed AIMessage) — DeepAgents often sends these - # as a full message rather than chunks, so tool_call_chunks would be empty. - for tc in getattr(chunk, "tool_calls", []) or []: - tc_id = _field(tc, "id") or "" - name = _field(tc, "name") or "" - args = _field(tc, "args") or {} - if tc_id and name: - pending_calls[tc_id] = { - "tool": name, - "args": args if isinstance(args, dict) else {}, - } - - - content = getattr(chunk, "content", "") - if content and isinstance(content, str): - if not getattr(chunk, "tool_call_id", None): - text_buf += content - new_clean = _clean(text_buf) - delta = new_clean[len(clean_buf):] - if delta: - clean_buf = new_clean - response_text += delta - yield f"data: {json.dumps({'type': 'token', 'text': delta})}\n\n" - - tc_id = getattr(chunk, "tool_call_id", None) - if tc_id: - _flush_text_buf() - - for entry in tc_accum.values(): - eid = entry["id"] - if eid: - try: - eargs: Any = json.loads(entry["args_str"]) if entry["args_str"] else {} - except json.JSONDecodeError: - eargs = {} - pending_calls[eid] = {"tool": entry["name"], "args": eargs} - tc_accum.clear() - - call_info = pending_calls.pop( - tc_id, - {"tool": getattr(chunk, "name", None) or "unknown", "args": {}}, - ) - try: - result = json.loads(content) - except (json.JSONDecodeError, TypeError): - result = {"raw": str(content)[:500]} - - logger.info( - " [{sid}] 🔧 {tool}({args})", - sid=sid, tool=call_info["tool"], args=call_info["args"], - ) - if isinstance(result, dict) and "error" in result: - logger.warning(" [{sid}] ⚠ error: {err}", sid=sid, err=result["error"]) - else: - keys = list(result.keys()) if isinstance(result, dict) else type(result).__name__ - count = result.get("count", result.get("events", result.get("functions", ""))) - logger.info( - " [{sid}] ✓ keys={keys} count={count}", - sid=sid, keys=keys, count=count if isinstance(count, int) else "—", + _flush_text_buf() + + for entry in tc_accum.values(): + eid = entry["id"] + if eid: + try: + eargs: Any = json.loads(entry["args_str"]) if entry["args_str"] else {} + except json.JSONDecodeError: + eargs = {} + pending_calls[eid] = {"tool": entry["name"], "args": eargs} + tc_accum.clear() + + call_info = pending_calls.pop( + tc_id, + {"tool": getattr(chunk, "name", None) or "unknown", "args": {}}, ) + try: + result = json.loads(content) + except (json.JSONDecodeError, TypeError): + result = {"raw": str(content)[:500]} - tool_calls_log.append({"tool": call_info["tool"], "args": call_info["args"], "result": result}) - label = _pick_label(call_info["tool"]) - yield f"data: {json.dumps({'type': 'tool_status', 'label': label})}\n\n" - yield f"data: {json.dumps({'type': 'tool_call', 'tool': call_info['tool'], 'args': call_info['args'], 'result': result})}\n\n" + logger.info( + " [{sid}] 🔧 {tool}({args})", + sid=sid, tool=call_info["tool"], args=call_info["args"], + ) + if isinstance(result, dict) and "error" in result: + logger.warning(" [{sid}] ⚠ error: {err}", sid=sid, err=result["error"]) + else: + keys = list(result.keys()) if isinstance(result, dict) else type(result).__name__ + count = result.get("count", result.get("events", result.get("functions", ""))) + logger.info( + " [{sid}] ✓ keys={keys} count={count}", + sid=sid, keys=keys, count=count if isinstance(count, int) else "—", + ) + + tool_calls_log.append({"tool": call_info["tool"], "args": call_info["args"], "result": result}) + label = _pick_label(call_info["tool"]) + yield f"data: {json.dumps({'type': 'tool_status', 'label': label})}\n\n" + yield f"data: {json.dumps({'type': 'tool_call', 'tool': call_info['tool'], 'args': call_info['args'], 'result': result})}\n\n" + + except TimeoutError: + timed_out = True + logger.warning("[{sid}] timed out after {timeout}s", sid=sid, timeout=settings.investigation_timeout) + msg = ( + f"Investigation timed out after {settings.investigation_timeout}s. " + "Try a narrower prompt or increase INVESTIGATION_TIMEOUT." + ) + yield f"data: {json.dumps({'type': 'error', 'message': msg})}\n\n" except Exception as e: logger.error("[{sid}] stream error: {err}", sid=sid, err=e) @@ -198,6 +210,7 @@ def _flush_text_buf(): usage: dict[str, Any] = { "latency_ms": int((time.time() - start) * 1000), "model": settings.llm_model, + "timed_out": timed_out, } if usage_meta: usage["input_tokens"] = _field(usage_meta, "input_tokens", 0) or 0 diff --git a/src/cli/ask.py b/src/cli/ask.py index 34a153a..2b09c2f 100644 --- a/src/cli/ask.py +++ b/src/cli/ask.py @@ -1,7 +1,8 @@ +from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeoutError from typing import Annotated import typer -from langchain_openai import ChatOpenAI +from langchain_litellm import ChatLiteLLM from langchain_core.messages import HumanMessage, SystemMessage from rich.console import Console from rich.markdown import Markdown @@ -16,10 +17,10 @@ def ask_cmd( question: Annotated[str, typer.Argument(help="Freeform question about your AWS environment.")], ) -> None: """Ask a freeform question about your AWS environment.""" - model = ChatOpenAI( - model=settings.openrouter_model, - api_key=settings.openrouter_api_key, - base_url=settings.openrouter_base_url, + model = ChatLiteLLM( + model=settings.llm_model, + api_base=settings.llm_api_base or None, + api_key=settings.llm_api_key or None, ) messages = [ @@ -30,8 +31,16 @@ def ask_cmd( HumanMessage(content=question), ] - with Progress(SpinnerColumn(), TextColumn("[progress.description]{task.description}"), transient=True) as progress: - progress.add_task("Thinking...", total=None) - response = model.invoke(messages) + try: + with Progress(SpinnerColumn(), TextColumn("[progress.description]{task.description}"), transient=True) as progress: + progress.add_task("Thinking...", total=None) + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(model.invoke, messages) + response = future.result(timeout=settings.investigation_timeout) + except FuturesTimeoutError: + console.print( + f"[red]Request timed out after {settings.investigation_timeout}s. Try narrowing the question scope.[/red]" + ) + raise typer.Exit(code=1) console.print(Markdown(response.content)) diff --git a/src/cli/investigate.py b/src/cli/investigate.py index 9d0faeb..cddf3cc 100644 --- a/src/cli/investigate.py +++ b/src/cli/investigate.py @@ -33,7 +33,13 @@ def investigate_cmd( with Progress(SpinnerColumn(), TextColumn("[progress.description]{task.description}"), transient=True) as progress: progress.add_task("Investigating...", total=None) - result = agent.investigate(investigation) + try: + result = agent.investigate(investigation) + except TimeoutError: + console.print( + "[red]Investigation timed out. Try narrowing the scope or increase INVESTIGATION_TIMEOUT.[/red]" + ) + raise typer.Exit(code=1) if json_output: rprint(json.dumps(result.raw_json or result.model_dump(), indent=2)) diff --git a/src/mcp_server.py b/src/mcp_server.py index ee01d8c..8411d8c 100644 --- a/src/mcp_server.py +++ b/src/mcp_server.py @@ -46,7 +46,7 @@ async def _init() -> None: async def _run_agent(prompt: str) -> dict[str, Any]: """Run the agent, persist the turn to DB, notify Slack, and return the result.""" await _init() - from agent.core import get_agent + from agent.core import ainvoke_with_timeout from agent.turns import save_turn, notify_slack thread_id = str(uuid.uuid4()) @@ -56,10 +56,31 @@ async def _run_agent(prompt: str) -> dict[str, Any]: } import time start = time.time() - result = await get_agent().ainvoke( - {"messages": [{"role": "user", "content": prompt}]}, - config=config, - ) + try: + result = await ainvoke_with_timeout( + {"messages": [{"role": "user", "content": prompt}]}, + config, + ) + except TimeoutError: + latency_ms = int((time.time() - start) * 1000) + timeout_msg = ( + f"Investigation timed out after {settings.investigation_timeout}s. " + "Try narrowing the prompt or increase INVESTIGATION_TIMEOUT." + ) + await save_turn( + thread_id, + prompt, + timeout_msg, + [], + { + "model": settings.llm_model, + "latency_ms": latency_ms, + "input_tokens": 0, + "output_tokens": 0, + "timed_out": True, + }, + ) + return {"type": "text", "session_id": thread_id, "text": timeout_msg} # Collect tool calls and assistant text from message history tool_calls_log: list[dict[str, Any]] = [] diff --git a/src/models/sessions.py b/src/models/sessions.py index 49a28a1..40847cf 100644 --- a/src/models/sessions.py +++ b/src/models/sessions.py @@ -2,7 +2,7 @@ from __future__ import annotations -from pydantic import BaseModel +from pydantic import BaseModel, Field class SessionSummary(BaseModel): @@ -10,12 +10,13 @@ class SessionSummary(BaseModel): title: str | None last_active_at: str | None model: str + aws_region: str class ToolCallRecord(BaseModel): tool_name: str args: dict - result: dict + result: dict | None error: str | None @@ -28,8 +29,9 @@ class UsageRecord(BaseModel): class MessageRecord(BaseModel): + id: str role: str content: str created_at: str | None - tool_calls: list[ToolCallRecord] = [] + tool_calls: list[ToolCallRecord] = Field(default_factory=list) usage: UsageRecord | None = None diff --git a/tests/conftest.py b/tests/conftest.py index b6bc57d..a6ce713 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -11,3 +11,13 @@ def aws_credentials(): os.environ.setdefault("AWS_SECURITY_TOKEN", "testing") os.environ.setdefault("AWS_SESSION_TOKEN", "testing") os.environ.setdefault("AWS_DEFAULT_REGION", "us-east-1") + + +@pytest.fixture(autouse=True) +def clear_tool_cache(): + """Ensure tool-level TTL cache never leaks state between tests.""" + from tools._cache import _cache + + _cache.clear() + yield + _cache.clear() diff --git a/tests/integration/test_poller.py b/tests/integration/test_poller.py index 781e73d..000400d 100644 --- a/tests/integration/test_poller.py +++ b/tests/integration/test_poller.py @@ -1,74 +1,100 @@ -"""Integration test for the proactive poller. +"""Assertive tests for proactive poller dispatch behavior.""" -Tests the full dispatch flow: fake alarm detected → investigation → Slack post. -Patches AWS data and the agent investigation so no real infra is needed. -""" +from __future__ import annotations import asyncio -import sys -from unittest.mock import patch, AsyncMock - -if sys.platform == "win32": - asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy()) - -FAKE_ALARM = { - "alarms": [ - { - "name": "HighErrorRate-payment-processor", - "reason": "Threshold Crossed: 3 datapoints were greater than threshold (5.0). Most recent: [18.4, 21.2, 19.8].", - "metric": "Errors", + +from agent import poller + + +def test_should_investigate_respects_cooldown(): + poller._last_investigated.clear() + key = "alarm:HighErrorRate" + + assert poller._should_investigate(key) is True + poller._mark_investigated(key) + assert poller._should_investigate(key) is False + + +def test_check_alarms_dispatches_once(monkeypatch): + poller._last_investigated.clear() + calls = {"run": 0, "persist": 0} + + def fake_get_alarms(_state): + return { + "alarms": [ + { + "name": "HighErrorRate-payment-processor", + "reason": "Threshold crossed", + "metric": "Errors", + } + ] + } + + async def fake_run(prompt, trigger_key): + calls["run"] += 1 + assert "HighErrorRate-payment-processor" in prompt + assert trigger_key == "alarm:HighErrorRate-payment-processor" + return { + "root_cause_category": "RESOURCE_LIMIT", + "root_cause_summary": "concurrency exhausted", + "confidence": "HIGH", + "evidence": ["throttles spiked"], + "mitigation_steps": ["increase reserved concurrency"], + "validation_steps": ["watch metrics"], + "services_affected": ["Lambda"], + "recommended_follow_up": "add alarm", } - ] -} - -FAKE_INVESTIGATION_RESULT = { - "root_cause_category": "RESOURCE_LIMIT", - "root_cause_summary": ( - "Lambda function payment-processor exceeded its concurrency limit during a traffic spike, " - "causing invocations to be throttled and error rate to climb above 18%." - ), - "confidence": "HIGH", - "evidence": [ - "CloudWatch alarm HighErrorRate-payment-processor crossed threshold at 18.4 errors/min", - "Concurrent executions metric shows saturation at reserved limit of 50", - "No recent deployments found in CloudTrail — config unchanged", - ], - "mitigation_steps": [ - "Increase reserved concurrency limit from 50 to 150 for payment-processor", - "Add SQS queue in front of the Lambda to absorb traffic spikes", - "Set a CloudWatch alarm on ConcurrentExecutions at 80% of reserved limit", - ], - "validation_steps": [ - "Confirm error rate drops below 1% after concurrency increase", - "Monitor throttle count metric for 15 minutes post-change", - ], - "services_affected": ["Lambda", "CloudWatch"], - "recommended_follow_up": "Review traffic patterns over the last 7 days to right-size concurrency limit.", -} - - -async def main(): - from agent.config import settings - from agent.db import db - from agent.core import init_agent - - if not settings.slack_webhook_url: - print("SLACK_WEBHOOK_URL not set — Slack post will be skipped.") - - print("Initialising agent…") - checkpointer = await db.init() - init_agent(checkpointer) - - print("Running _check_alarms() with a fake alarm + fake investigation…") - with ( - patch("tools.cloudwatch.get_alarms", return_value=FAKE_ALARM), - patch("agent.poller._run_investigation", new=AsyncMock(return_value=FAKE_INVESTIGATION_RESULT)), - ): - from agent.poller import _check_alarms - await _check_alarms() - - print("Done. Check your Slack channel.") - await db.close() - - -asyncio.run(main()) + + async def fake_persist(_prompt, _result, _session_id): + calls["persist"] += 1 + + monkeypatch.setattr("tools.cloudwatch.get_alarms", fake_get_alarms) + monkeypatch.setattr(poller, "_run_investigation", fake_run) + monkeypatch.setattr(poller, "_persist_and_notify", fake_persist) + + asyncio.run(poller._check_alarms()) + asyncio.run(poller._check_alarms()) + + assert calls["run"] == 1 + assert calls["persist"] == 1 + + +def test_check_lambda_errors_dispatches_for_threshold_breach(monkeypatch): + poller._last_investigated.clear() + calls = {"run": 0, "persist": 0} + + def fake_list_lambda_functions(): + return {"functions": [{"name": "payment-processor"}]} + + def fake_get_lambda_error_rate(_name, _hours): + return {"error_rate_pct": 12.5} + + async def fake_run(prompt, trigger_key): + calls["run"] += 1 + assert "payment-processor" in prompt + assert trigger_key == "lambda_errors:payment-processor" + return { + "root_cause_category": "RESOURCE_LIMIT", + "root_cause_summary": "lambda saturation", + "confidence": "HIGH", + "evidence": ["high throttles"], + "mitigation_steps": ["raise concurrency"], + "validation_steps": ["observe errors"], + "services_affected": ["Lambda"], + "recommended_follow_up": "right-size limits", + } + + async def fake_persist(_prompt, _result, _session_id): + calls["persist"] += 1 + + monkeypatch.setattr("tools.lambda_.list_lambda_functions", fake_list_lambda_functions) + monkeypatch.setattr("tools.lambda_.get_lambda_error_rate", fake_get_lambda_error_rate) + monkeypatch.setattr(poller, "_run_investigation", fake_run) + monkeypatch.setattr(poller, "_persist_and_notify", fake_persist) + monkeypatch.setattr(poller.settings, "poll_error_threshold", 5.0, raising=False) + + asyncio.run(poller._check_lambda_errors()) + + assert calls["run"] == 1 + assert calls["persist"] == 1 diff --git a/tests/integration/test_slack.py b/tests/integration/test_slack.py index 3e85492..5482a6b 100644 --- a/tests/integration/test_slack.py +++ b/tests/integration/test_slack.py @@ -1,36 +1,69 @@ -"""Quick standalone test — sends a sample investigation result to Slack.""" +"""Assertive tests for Slack payload rendering and webhook delivery.""" + +from __future__ import annotations import asyncio -from agent.config import settings -from integrations.slack_webhook import post_investigation + +from integrations.slack_webhook import _build_payload, post_investigation SAMPLE_RESULT = { "root_cause_category": "RESOURCE_LIMIT", - "root_cause_summary": ( - "Lambda function payment-processor exceeded its memory limit (512 MB) " - "during a traffic spike, causing 47 invocations to fail with OOM errors." - ), + "root_cause_summary": "Lambda exhausted memory under burst traffic.", "confidence": "HIGH", "evidence": [ - "CloudWatch: MemoryUtilization hit 99.8% at 14:32 UTC", - "Lambda logs: 47x 'Runtime exited with error: signal: killed' in the last hour", - "Error rate jumped from 0.2% → 18.4% at 14:31 UTC", + "MemoryUtilization reached 99%", + "OOM errors spiked in CloudWatch logs", ], "mitigation_steps": [ - "Increase Lambda memory from 512 MB to 1024 MB", - "Add a memory utilization alarm at 80% to catch this earlier", - "Review payload sizes — large S3 objects may be loaded into memory unnecessarily", + "Increase Lambda memory", + "Add memory alarm at 80%", ], "services_affected": ["Lambda", "CloudWatch"], } -async def main(): - if not settings.slack_webhook_url: - print("SLACK_WEBHOOK_URL is not set in .env — aborting.") - return - print(f"Sending test message to Slack webhook …") - await post_investigation(settings.slack_webhook_url, SAMPLE_RESULT, "test-session-1234") - print("Done.") +def test_build_payload_contains_category_confidence_and_summary(): + payload = _build_payload(SAMPLE_RESULT, "session-1234") + attachment = payload["attachments"][0] + blocks = attachment["blocks"] + + assert attachment["color"] == "#e74c3c" + assert any("RESOURCE_LIMIT" in str(block) for block in blocks) + assert any("HIGH" in str(block) for block in blocks) + assert any("Lambda exhausted memory" in str(block) for block in blocks) + + +def test_post_investigation_sends_payload(monkeypatch): + captured = {} + + class FakeResponse: + status_code = 200 + text = "ok" + + class FakeClient: + def __init__(self, timeout): + self.timeout = timeout + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + _ = (exc_type, exc, tb) + + async def post(self, url, json): + captured["url"] = url + captured["json"] = json + return FakeResponse() + + monkeypatch.setattr("httpx.AsyncClient", FakeClient) + + asyncio.run( + post_investigation( + "https://hooks.slack.test/example", + SAMPLE_RESULT, + "session-1234", + ) + ) -asyncio.run(main()) + assert captured["url"] == "https://hooks.slack.test/example" + assert "attachments" in captured["json"] diff --git a/tests/test_agent/test_core.py b/tests/test_agent/test_core.py new file mode 100644 index 0000000..a0380de --- /dev/null +++ b/tests/test_agent/test_core.py @@ -0,0 +1,75 @@ +"""Agent core tests: completion parsing, timeout behavior, and failure paths.""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import pytest + +from agent import core +from models.agent import Investigation + + +class _FastAgent: + async def ainvoke(self, _payload, config=None): + return {"messages": [SimpleNamespace(content="ok")], "config": config} + + +class _SlowAgent: + async def ainvoke(self, _payload, config=None): + _ = config + await asyncio.sleep(0.2) + return {"messages": [SimpleNamespace(content="late")]} + + +def test_ainvoke_with_timeout_success(monkeypatch): + monkeypatch.setattr(core, "_agent", _FastAgent()) + monkeypatch.setattr(core.settings, "investigation_timeout", 1) + + result = asyncio.run(core.ainvoke_with_timeout({"messages": []}, {"configurable": {}})) + assert result["messages"][0].content == "ok" + + +def test_ainvoke_with_timeout_raises_timeout(monkeypatch): + monkeypatch.setattr(core, "_agent", _SlowAgent()) + monkeypatch.setattr(core.settings, "investigation_timeout", 0.01) + + with pytest.raises(TimeoutError): + asyncio.run(core.ainvoke_with_timeout({"messages": []}, {"configurable": {}})) + + +def test_investigate_completion_path_parses_structured_json(monkeypatch): + raw_json = { + "root_cause_category": "RESOURCE_LIMIT", + "root_cause_summary": "Lambda concurrency exhausted", + "evidence": ["Throttles spiked"], + "mitigation_steps": ["Increase reserved concurrency"], + "validation_steps": ["Watch throttles for 15m"], + "confidence": "HIGH", + "services_affected": ["Lambda"], + "recommended_follow_up": "Add alarm on concurrency", + } + payload = "```json\n" + __import__("json").dumps(raw_json) + "\n```" + + def fake_invoke(_input, _config): + return {"messages": [SimpleNamespace(content=payload)]} + + monkeypatch.setattr(core, "ensure_agent_initialized", lambda: None) + monkeypatch.setattr(core, "invoke_with_timeout", fake_invoke) + + result = core.InvestigationAgent().investigate(Investigation(description="high Lambda errors")) + assert result.root_cause_summary == "Lambda concurrency exhausted" + assert result.confidence.value == "HIGH" + assert result.root_cause_category.value == "RESOURCE_LIMIT" + + +def test_investigate_failure_path_propagates_agent_error(monkeypatch): + def failing_invoke(_input, _config): + raise RuntimeError("recursion limit reached") + + monkeypatch.setattr(core, "ensure_agent_initialized", lambda: None) + monkeypatch.setattr(core, "invoke_with_timeout", failing_invoke) + + with pytest.raises(RuntimeError, match="recursion"): + core.InvestigationAgent().investigate(Investigation(description="check recursion failure")) diff --git a/tests/test_tools/test_bash_tool.py b/tests/test_tools/test_bash_tool.py index 0686c88..0a776d1 100644 --- a/tests/test_tools/test_bash_tool.py +++ b/tests/test_tools/test_bash_tool.py @@ -9,7 +9,6 @@ from tools.bash_tool import run_bash_command - # ── Helpers ────────────────────────────────────────────────────────────────── def _ok_proc(stdout: str = "output", stderr: str = "") -> MagicMock: diff --git a/tests/test_tools/test_cloudwatch.py b/tests/test_tools/test_cloudwatch.py index 01275f9..2c0ba16 100644 --- a/tests/test_tools/test_cloudwatch.py +++ b/tests/test_tools/test_cloudwatch.py @@ -1,20 +1,19 @@ import boto3 -import pytest from moto import mock_aws from tools.cloudwatch import ( - DescribeLogGroupsTool, - GetAlarmHistoryTool, - GetAlarmsTool, - GetLogEventsTool, - GetMetricDataTool, + describe_log_groups, + get_alarm_history, + get_alarms, + get_log_events, + get_metric_data, + query_logs_insights, ) @mock_aws def test_get_alarms_empty(): - tool = GetAlarmsTool() - result = tool.run() + result = get_alarms() assert result["alarms"] == [] assert result["count"] == 0 @@ -32,8 +31,7 @@ def test_get_alarms_with_alarm(): Threshold=1, ComparisonOperator="GreaterThanThreshold", ) - tool = GetAlarmsTool() - result = tool.run() + result = get_alarms() assert result["count"] == 1 assert result["alarms"][0]["name"] == "test-alarm" @@ -51,8 +49,7 @@ def test_get_alarms_filter_by_state(): Threshold=1, ComparisonOperator="GreaterThanThreshold", ) - tool = GetAlarmsTool() - result = tool.run(state="ALARM") + result = get_alarms(state="ALARM") assert isinstance(result["alarms"], list) @@ -69,16 +66,14 @@ def test_get_alarm_history(): Threshold=1, ComparisonOperator="GreaterThanThreshold", ) - tool = GetAlarmHistoryTool() - result = tool.run(alarm_name="history-alarm", hours=24) + result = get_alarm_history(alarm_name="history-alarm", hours=24) # moto may not implement describe_alarm_history; we accept either a result or a graceful error assert "history" in result or "error" in result @mock_aws def test_get_metric_data(): - tool = GetMetricDataTool() - result = tool.run( + result = get_metric_data( namespace="AWS/Lambda", metric="Errors", dimensions=[{"Name": "FunctionName", "Value": "my-fn"}], @@ -89,8 +84,7 @@ def test_get_metric_data(): @mock_aws def test_describe_log_groups_empty(): - tool = DescribeLogGroupsTool() - result = tool.run() + result = describe_log_groups() assert result["log_groups"] == [] assert result["count"] == 0 @@ -99,8 +93,7 @@ def test_describe_log_groups_empty(): def test_describe_log_groups_with_group(): client = boto3.client("logs", region_name="us-east-1") client.create_log_group(logGroupName="/aws/lambda/my-fn") - tool = DescribeLogGroupsTool() - result = tool.run() + result = describe_log_groups() assert result["count"] == 1 assert result["log_groups"][0]["name"] == "/aws/lambda/my-fn" @@ -109,6 +102,17 @@ def test_describe_log_groups_with_group(): def test_get_log_events_empty(): client = boto3.client("logs", region_name="us-east-1") client.create_log_group(logGroupName="/aws/lambda/my-fn") - tool = GetLogEventsTool() - result = tool.run(log_group="/aws/lambda/my-fn") + result = get_log_events(log_group="/aws/lambda/my-fn") assert result["events"] == [] + + +@mock_aws +def test_query_logs_insights_graceful_result(): + client = boto3.client("logs", region_name="us-east-1") + client.create_log_group(logGroupName="/aws/lambda/my-fn") + result = query_logs_insights( + log_group="/aws/lambda/my-fn", + query="fields @timestamp, @message | limit 5", + hours=1, + ) + assert "results" in result or "error" in result diff --git a/tests/test_tools/test_ecs.py b/tests/test_tools/test_ecs.py index b4c8a82..41ea9a0 100644 --- a/tests/test_tools/test_ecs.py +++ b/tests/test_tools/test_ecs.py @@ -1,16 +1,30 @@ import boto3 -import pytest from moto import mock_aws -from tools.ecs import DescribeServiceTool, ListServicesTool +from tools.ecs import describe_ecs_service, list_ecs_clusters, list_ecs_services + + +@mock_aws +def test_list_clusters_empty(): + result = list_ecs_clusters() + assert result["clusters"] == [] + assert result["count"] == 0 + + +@mock_aws +def test_list_clusters_with_cluster(): + client = boto3.client("ecs", region_name="us-east-1") + client.create_cluster(clusterName="test-cluster") + result = list_ecs_clusters() + assert result["count"] == 1 + assert result["clusters"][0]["name"] == "test-cluster" @mock_aws def test_list_services_empty_cluster(): client = boto3.client("ecs", region_name="us-east-1") client.create_cluster(clusterName="test-cluster") - tool = ListServicesTool() - result = tool.run(cluster="test-cluster") + result = list_ecs_services(cluster="test-cluster") assert result["services"] == [] assert result["count"] == 0 @@ -19,6 +33,5 @@ def test_list_services_empty_cluster(): def test_describe_service_not_found(): client = boto3.client("ecs", region_name="us-east-1") client.create_cluster(clusterName="test-cluster") - tool = DescribeServiceTool() - result = tool.run(cluster="test-cluster", service="nonexistent") + result = describe_ecs_service(cluster="test-cluster", service="nonexistent") assert "error" in result diff --git a/tests/test_tools/test_lambda.py b/tests/test_tools/test_lambda.py index d565ddb..7fced2c 100644 --- a/tests/test_tools/test_lambda.py +++ b/tests/test_tools/test_lambda.py @@ -1,8 +1,7 @@ import boto3 -import pytest from moto import mock_aws -from tools.lambda_ import GetFunctionConfigTool, GetLambdaErrorRateTool, ListFunctionsTool +from tools.lambda_ import get_lambda_error_rate, get_lambda_function_config, list_lambda_functions def _create_lambda_role(iam_client: boto3.client) -> str: @@ -27,8 +26,7 @@ def _create_lambda_role(iam_client: boto3.client) -> str: @mock_aws def test_list_functions_empty(): - tool = ListFunctionsTool() - result = tool.run() + result = list_lambda_functions() assert result["functions"] == [] assert result["count"] == 0 @@ -47,8 +45,7 @@ def test_list_functions_with_function(): MemorySize=256, Timeout=30, ) - tool = ListFunctionsTool() - result = tool.run() + result = list_lambda_functions() assert result["count"] == 1 fn = result["functions"][0] assert fn["name"] == "my-fn" @@ -69,14 +66,12 @@ def test_get_function_config(): Code={"ZipFile": b"fake"}, Environment={"Variables": {"KEY1": "val1", "KEY2": "val2"}}, ) - tool = GetFunctionConfigTool() - result = tool.run(name="cfg-fn") + result = get_lambda_function_config(name="cfg-fn") assert result["name"] == "cfg-fn" assert set(result["env_var_keys"]) == {"KEY1", "KEY2"} @mock_aws def test_get_lambda_error_rate(): - tool = GetLambdaErrorRateTool() - result = tool.run(name="nonexistent-fn", hours=1) + result = get_lambda_error_rate(name="nonexistent-fn", hours=1) assert "total_errors" in result or "error" in result diff --git a/tests/test_tools/test_rds.py b/tests/test_tools/test_rds.py index 2345076..d5b93ab 100644 --- a/tests/test_tools/test_rds.py +++ b/tests/test_tools/test_rds.py @@ -1,14 +1,12 @@ import boto3 -import pytest from moto import mock_aws -from tools.rds import DescribeDBInstancesTool, GetDBEventsTool +from tools.rds import describe_rds_instances, get_rds_events @mock_aws def test_describe_rds_instances_empty(): - tool = DescribeDBInstancesTool() - result = tool.run() + result = describe_rds_instances() assert result["instances"] == [] assert result["count"] == 0 @@ -24,14 +22,12 @@ def test_describe_rds_instances_with_db(): MasterUserPassword="password123", AllocatedStorage=20, ) - tool = DescribeDBInstancesTool() - result = tool.run() + result = describe_rds_instances() assert result["count"] == 1 assert result["instances"][0]["identifier"] == "mydb" @mock_aws def test_get_rds_events_empty(): - tool = GetDBEventsTool() - result = tool.run(hours=1) + result = get_rds_events(hours=1) assert "events" in result From 7c2d1ae09a9d7828064ea925d942b2cc2357d876 Mon Sep 17 00:00:00 2001 From: AhmadHammad21 Date: Tue, 5 May 2026 22:27:00 +0300 Subject: [PATCH 6/6] feat: harden bash tool and docker runtime --- Dockerfile | 15 +++- docker-compose.yml | 32 ++++++-- frontend/Dockerfile | 13 +++- frontend/nginx.conf | 2 +- scripts/test_db_connection.py | 109 ++++++++++++++++++++++++++ src/tools/bash_tool.py | 118 +++++++++++++++++++++++++---- tests/test_tools/test_bash_tool.py | 10 +++ 7 files changed, 270 insertions(+), 29 deletions(-) create mode 100644 scripts/test_db_connection.py diff --git a/Dockerfile b/Dockerfile index 201c7e3..bc194dc 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,14 +1,14 @@ # ── Stage 1: install Python dependencies with uv ───────────────────────────── -FROM python:3.12-slim AS deps +FROM python:3.12.9-slim-bookworm AS deps -COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/ +COPY --from=ghcr.io/astral-sh/uv:0.6.17 /uv /uvx /bin/ WORKDIR /app COPY pyproject.toml uv.lock ./ RUN uv sync --frozen --no-dev --no-install-project # ── Stage 2: final image ────────────────────────────────────────────────────── -FROM python:3.12-slim +FROM python:3.12.9-slim-bookworm # AWS CLI v2 — required for the bash execution tool (run_bash_command) RUN apt-get update && apt-get install -y --no-install-recommends \ @@ -21,7 +21,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ && apt-get clean \ && rm -rf /var/lib/apt/lists/* -COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/ +COPY --from=ghcr.io/astral-sh/uv:0.6.17 /uv /uvx /bin/ WORKDIR /app @@ -33,10 +33,17 @@ COPY src/ src/ COPY migrations/ migrations/ COPY scripts/ scripts/ +RUN groupadd --gid 10001 app && useradd --uid 10001 --gid app --create-home --shell /usr/sbin/nologin app +RUN chown -R app:app /app + ENV PATH="/app/.venv/bin:$PATH" \ PYTHONPATH="/app/src" \ + HOME="/home/app" \ + PYTHONDONTWRITEBYTECODE=1 \ PYTHONUNBUFFERED=1 +USER app + EXPOSE 8000 CMD ["uvicorn", "api.app:app", "--host", "0.0.0.0", "--port", "8000"] diff --git a/docker-compose.yml b/docker-compose.yml index 9a119fb..c940ffe 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -2,13 +2,13 @@ services: # ── PostgreSQL ────────────────────────────────────────────────────────────── db: - image: postgres:16 + image: postgres:16.4 environment: POSTGRES_DB: opendevops POSTGRES_USER: dev POSTGRES_PASSWORD: dev ports: - - "5432:5432" + - "5433:5432" volumes: - postgres_data:/var/lib/postgresql/data healthcheck: @@ -16,6 +16,7 @@ services: interval: 5s timeout: 5s retries: 5 + restart: unless-stopped # ── FastAPI backend ───────────────────────────────────────────────────────── backend: @@ -24,23 +25,44 @@ services: - "8000:8000" env_file: .env environment: - # Override DATABASE_URL to point at the compose db service + # Keep one canonical variable name; container uses docker-internal host. DATABASE_URL: postgresql://dev:dev@db:5432/opendevops volumes: # Mount host AWS credentials read-only so the bash tool can call the CLI. # For production on AWS (ECS/EC2), remove this and use an IAM role instead. - - ${HOME}/.aws:/root/.aws:ro + - ${USERPROFILE}/.aws:/home/app/.aws:ro depends_on: db: condition: service_healthy + security_opt: + - no-new-privileges:true + cap_drop: + - ALL + read_only: true + tmpfs: + - /tmp + healthcheck: + test: ["CMD", "python", "-c", "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8000/openapi.json', timeout=5)"] + interval: 15s + timeout: 5s + retries: 5 + restart: unless-stopped # ── React frontend (production build served by nginx) ────────────────────── frontend: build: ./frontend ports: - - "80:80" + - "80:8080" depends_on: - backend + security_opt: + - no-new-privileges:true + healthcheck: + test: ["CMD-SHELL", "wget -q -O /dev/null http://127.0.0.1:8080/ || exit 1"] + interval: 15s + timeout: 5s + retries: 5 + restart: unless-stopped volumes: postgres_data: diff --git a/frontend/Dockerfile b/frontend/Dockerfile index 3b0ed25..7fe1ef8 100644 --- a/frontend/Dockerfile +++ b/frontend/Dockerfile @@ -1,5 +1,5 @@ # Production build — React app served by nginx -FROM node:20-slim AS build +FROM node:20-bookworm-slim AS build WORKDIR /app COPY package.json package-lock.json* ./ @@ -7,7 +7,14 @@ RUN npm ci COPY . . RUN npm run build -FROM nginx:alpine +FROM nginx:1.27-alpine COPY --from=build /app/dist /usr/share/nginx/html COPY nginx.conf /etc/nginx/conf.d/default.conf -EXPOSE 80 +RUN mkdir -p /var/cache/nginx /var/run /var/log/nginx /tmp/nginx \ + && chown -R nginx:nginx /usr/share/nginx/html /etc/nginx/conf.d /var/cache/nginx /var/run /var/log/nginx /tmp/nginx \ + && sed -ri 's@^\s*pid\s+[^;]+;@pid /tmp/nginx.pid;@' /etc/nginx/nginx.conf + +USER nginx +EXPOSE 8080 + +CMD ["nginx", "-g", "daemon off;"] diff --git a/frontend/nginx.conf b/frontend/nginx.conf index 1b36089..2549a9e 100644 --- a/frontend/nginx.conf +++ b/frontend/nginx.conf @@ -1,5 +1,5 @@ server { - listen 80; + listen 8080; root /usr/share/nginx/html; index index.html; diff --git a/scripts/test_db_connection.py b/scripts/test_db_connection.py new file mode 100644 index 0000000..b8a5472 --- /dev/null +++ b/scripts/test_db_connection.py @@ -0,0 +1,109 @@ +""" +Database connectivity smoke test. + +Reads .env, connects to PostgreSQL, and lists user tables. + +Usage: + uv run python scripts/test_db_connection.py + uv run python scripts/test_db_connection.py --url-var DOCKER_DATABASE_URL +""" + +from __future__ import annotations + +import argparse +import os +import sys +from pathlib import Path +from urllib.parse import urlparse, urlunparse + +from dotenv import load_dotenv + + +def _mask_url(url: str) -> str: + try: + parsed = urlparse(url) + if parsed.password: + netloc = parsed.netloc.replace(parsed.password, "***") + return urlunparse(parsed._replace(netloc=netloc)) + return url + except Exception: + return "" + + +def _print(msg: str) -> None: + print(msg) + + +def main() -> int: + parser = argparse.ArgumentParser(description="Test PostgreSQL connectivity and list tables") + parser.add_argument( + "--url-var", + default="DATABASE_URL", + choices=["DATABASE_URL", "DOCKER_DATABASE_URL"], + help="Environment variable to use from .env", + ) + args = parser.parse_args() + + repo_root = Path(__file__).parent.parent + load_dotenv(repo_root / ".env") + + db_url = os.getenv(args.url_var) + if not db_url: + _print(f"ERROR: {args.url_var} is not set in .env") + _print("Hint: set DATABASE_URL for local host access, and DOCKER_DATABASE_URL for docker-internal access") + return 1 + + _print(f"Using {args.url_var}: {_mask_url(db_url)}") + + try: + import psycopg # type: ignore + + with psycopg.connect(db_url, connect_timeout=5) as conn: + with conn.cursor() as cur: + cur.execute( + """ + SELECT + current_database(), + current_user, + inet_server_addr()::text, + inet_server_port() + """ + ) + db_name, db_user, server_addr, server_port = cur.fetchone() + _print( + "Connected: " + f"database={db_name}, user={db_user}, server={server_addr}:{server_port}" + ) + + cur.execute( + """ + SELECT table_schema, table_name + FROM information_schema.tables + WHERE table_type = 'BASE TABLE' + AND table_schema NOT IN ('pg_catalog', 'information_schema') + ORDER BY table_schema, table_name + """ + ) + rows = cur.fetchall() + + if not rows: + _print("No user tables found.") + return 0 + + _print(f"Found {len(rows)} table(s):") + for schema, table in rows: + _print(f" - {schema}.{table}") + return 0 + + except Exception as exc: + _print(f"ERROR: Could not connect/query database: {exc}") + if args.url_var == "DOCKER_DATABASE_URL": + _print("Hint: DOCKER_DATABASE_URL usually works only inside Docker network.") + _print("Run with DATABASE_URL from host, or run this script from docker compose backend service.") + else: + _print("Hint: ensure docker compose db is running and port 5432 is available on host.") + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/tools/bash_tool.py b/src/tools/bash_tool.py index 81337ba..957b676 100644 --- a/src/tools/bash_tool.py +++ b/src/tools/bash_tool.py @@ -22,27 +22,109 @@ "batch-get", }) -# kubectl and docker use exact prefix matching (narrower — fewer read-only verbs to reason about) -_KUBECTL_PREFIXES = ("kubectl get", "kubectl describe", "kubectl logs") -_DOCKER_PREFIXES = ("docker ps", "docker logs", "docker inspect") +_COMMAND_MAX_LEN = 2000 +_CHAIN_TOKENS: frozenset[str] = frozenset({";", "&&", "||", "|", ">", ">>", "<"}) + +_AWS_GLOBAL_FLAGS_WITH_VALUE: frozenset[str] = frozenset({ + "--profile", + "--region", + "--output", + "--query", + "--endpoint-url", + "--ca-bundle", + "--cli-connect-timeout", + "--cli-read-timeout", +}) +_AWS_GLOBAL_FLAGS_NO_VALUE: frozenset[str] = frozenset({ + "--debug", + "--no-cli-pager", + "--no-verify-ssl", + "--color", +}) +_AWS_BLOCKED_GLOBAL_FLAGS: frozenset[str] = frozenset({"--endpoint-url"}) + +_KUBECTL_FLAGS_WITH_VALUE: frozenset[str] = frozenset({ + "-n", + "--namespace", + "--context", + "--cluster", + "--user", + "--kubeconfig", + "--server", + "--request-timeout", +}) + +_DOCKER_FLAGS_WITH_VALUE: frozenset[str] = frozenset({"-H", "--host", "--context", "--config"}) _TIMEOUT = 30 _MAX_OUTPUT = 4000 _MAX_STDERR = 1000 -def _allowed(command: str) -> bool: - parts = command.split() - if not parts: +def _tokenize(command: str) -> list[str]: + try: + return shlex.split(command) + except ValueError: + return [] + + +def _find_subcommand( + tokens: list[str], + start_idx: int, + flags_with_value: frozenset[str], +) -> str | None: + i = start_idx + while i < len(tokens): + tok = tokens[i] + if tok.startswith("-"): + if tok in flags_with_value: + i += 2 + continue + i += 1 + continue + return tok.lower() + return None + + +def _allowed(command: str, tokens: list[str]) -> bool: + if not tokens: + return False + if len(command) > _COMMAND_MAX_LEN: + return False + if any(c in command for c in ("\n", "\r", "\x00")): + return False + if any(tok in _CHAIN_TOKENS for tok in tokens): + return False + if any(marker in command for marker in ("&&", "||", ";", "|", ">", "<")): return False - binary = parts[0] + binary = tokens[0] if binary == "aws": - # Require: aws — at least 3 tokens - if len(parts) < 3: + # Allow global flags before service/op: aws [flags] + i = 1 + while i < len(tokens) and tokens[i].startswith("-"): + flag = tokens[i].lower() + if flag in _AWS_BLOCKED_GLOBAL_FLAGS: + return False + if flag in _AWS_GLOBAL_FLAGS_WITH_VALUE: + if i + 1 >= len(tokens): + return False + i += 2 + continue + if flag in _AWS_GLOBAL_FLAGS_NO_VALUE: + i += 1 + continue + # Unknown top-level aws flags are rejected for safety. + return False + + # Require: aws [flags] + if i + 1 >= len(tokens): + return False + service = tokens[i] + operation = tokens[i + 1].lower() + if service.startswith("-"): return False - operation = parts[2].lower() # Extract verb = everything before the first dash (e.g. "describe-instances" → "describe") verb = operation.split("-")[0] # "batch-get-item" → verb fragment is "batch"; check "batch-get" prefix separately @@ -51,10 +133,12 @@ def _allowed(command: str) -> bool: return verb in _AWS_READONLY_VERBS if binary == "kubectl": - return any(command.startswith(p) for p in _KUBECTL_PREFIXES) + subcommand = _find_subcommand(tokens, 1, _KUBECTL_FLAGS_WITH_VALUE) + return subcommand in {"get", "describe", "logs"} if binary == "docker": - return any(command.startswith(p) for p in _DOCKER_PREFIXES) + subcommand = _find_subcommand(tokens, 1, _DOCKER_FLAGS_WITH_VALUE) + return subcommand in {"ps", "logs", "inspect"} return False @@ -78,16 +162,18 @@ def run_bash_command(command: str) -> dict[str, Any]: command: Read-only shell command to run. """ command = command.strip() + tokens = _tokenize(command) start = time.monotonic() - if not _allowed(command): + if not _allowed(command, tokens): logger.warning("bash_tool BLOCKED | cmd={!r}", command) return { "success": False, "output": "", "error": ( - "Command blocked. AWS commands must use a read-only operation verb " - "(describe-, list-, get-, lookup-, filter-, search-, scan-, query, batch-get-). " + "Command blocked. AWS commands must use safe global flags and a " + "read-only operation verb (describe-, list-, get-, lookup-, filter-, " + "search-, scan-, query, batch-get-). " "kubectl: only get/describe/logs. docker: only ps/logs/inspect." ), "command": command, @@ -97,7 +183,7 @@ def run_bash_command(command: str) -> dict[str, Any]: logger.info("bash_tool RUN | cmd={!r}", command) try: proc = subprocess.run( - shlex.split(command), + tokens, capture_output=True, text=True, timeout=_TIMEOUT, diff --git a/tests/test_tools/test_bash_tool.py b/tests/test_tools/test_bash_tool.py index 0a776d1..a20c45e 100644 --- a/tests/test_tools/test_bash_tool.py +++ b/tests/test_tools/test_bash_tool.py @@ -28,6 +28,7 @@ def _ok_proc(stdout: str = "output", stderr: str = "") -> MagicMock: "aws lambda get-function --function-name my-fn", "aws lambda list-functions", "aws ec2 describe-instances", + "aws --region us-east-1 ec2 describe-instances", "aws rds describe-db-instances", "aws cloudtrail lookup-events --max-results 10", # Extended AWS services — no boto3 tool defined for these @@ -56,10 +57,12 @@ def _ok_proc(stdout: str = "output", stderr: str = "") -> MagicMock: "aws stepfunctions list-state-machines", # kubectl and docker "kubectl get pods", + "kubectl -n kube-system get pods", "kubectl get pods -n kube-system", "kubectl describe pod my-pod", "kubectl logs my-pod", "docker ps", + "docker --context default ps", "docker logs my-container", "docker inspect my-container", ]) @@ -90,6 +93,7 @@ def test_allowlisted_commands_pass(mocker, cmd): "aws sns publish --topic-arn arn:aws:sns:us-east-1:123:my-topic --message hi", "aws sqs send-message --queue-url https://sqs.us-east-1.amazonaws.com/123/q", "aws rds modify-db-instance --db-instance-identifier my-db", + "aws --endpoint-url http://localhost:4566 s3api list-buckets", # kubectl write operations "kubectl delete pod my-pod", "kubectl apply -f deployment.yaml", @@ -98,6 +102,10 @@ def test_allowlisted_commands_pass(mocker, cmd): "docker rm my-container", "docker stop my-container", "docker run my-image", + # command chaining / token-smuggling attempts + "aws logs describe-log-groups && whoami", + "kubectl get pods; cat /etc/passwd", + "docker ps | cat", # arbitrary shell commands "rm -rf /", "bash -c 'echo pwned'", @@ -107,6 +115,8 @@ def test_allowlisted_commands_pass(mocker, cmd): # malformed / empty "aws", "aws s3", + "aws --region", + "aws logs describe-log-groups\nwhoami", "", " ", ])