codeshield-runtime 0.1.1__tar.gz → 0.1.3__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (21) hide show
  1. {codeshield_runtime-0.1.1 → codeshield_runtime-0.1.3}/PKG-INFO +59 -14
  2. {codeshield_runtime-0.1.1 → codeshield_runtime-0.1.3}/README.md +58 -13
  3. {codeshield_runtime-0.1.1 → codeshield_runtime-0.1.3}/pyproject.toml +1 -1
  4. codeshield_runtime-0.1.3/src/codeshield/__init__.py +46 -0
  5. {codeshield_runtime-0.1.1 → codeshield_runtime-0.1.3}/src/codeshield/environment.py +7 -1
  6. {codeshield_runtime-0.1.1 → codeshield_runtime-0.1.3}/src/codeshield/loop.py +24 -4
  7. {codeshield_runtime-0.1.1 → codeshield_runtime-0.1.3}/src/codeshield/runner.py +2 -1
  8. {codeshield_runtime-0.1.1 → codeshield_runtime-0.1.3}/src/codeshield/tools.py +13 -7
  9. {codeshield_runtime-0.1.1 → codeshield_runtime-0.1.3}/tests/test_coverage.py +7 -5
  10. {codeshield_runtime-0.1.1 → codeshield_runtime-0.1.3}/tests/test_engine.py +42 -6
  11. {codeshield_runtime-0.1.1 → codeshield_runtime-0.1.3}/tests/test_tools.py +10 -2
  12. codeshield_runtime-0.1.1/src/codeshield/__init__.py +0 -17
  13. {codeshield_runtime-0.1.1 → codeshield_runtime-0.1.3}/.gitignore +0 -0
  14. {codeshield_runtime-0.1.1 → codeshield_runtime-0.1.3}/LICENSE +0 -0
  15. {codeshield_runtime-0.1.1 → codeshield_runtime-0.1.3}/src/codeshield/__main__.py +0 -0
  16. {codeshield_runtime-0.1.1 → codeshield_runtime-0.1.3}/src/codeshield/analyzer.py +0 -0
  17. {codeshield_runtime-0.1.1 → codeshield_runtime-0.1.3}/src/codeshield/classifier.py +0 -0
  18. {codeshield_runtime-0.1.1 → codeshield_runtime-0.1.3}/src/codeshield/cli.py +0 -0
  19. {codeshield_runtime-0.1.1 → codeshield_runtime-0.1.3}/src/codeshield/schemas.py +0 -0
  20. {codeshield_runtime-0.1.1 → codeshield_runtime-0.1.3}/tests/__init__.py +0 -0
  21. {codeshield_runtime-0.1.1 → codeshield_runtime-0.1.3}/tests/test_cli.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: codeshield-runtime
3
- Version: 0.1.1
3
+ Version: 0.1.3
4
4
  Summary: A secure, isolated, and self-healing Python code execution engine powered by uv and AST analysis.
5
5
  Project-URL: Homepage, https://github.com/AlgorithmicMind/codeshield
6
6
  Project-URL: Repository, https://github.com/AlgorithmicMind/codeshield
@@ -40,6 +40,8 @@ Description-Content-Type: text/markdown
40
40
  [![Python](https://img.shields.io/badge/python-3.11%20%7C%203.12-blue)](https://www.python.org/)
41
41
  [![License: MIT](https://img.shields.io/badge/license-MIT-green)](LICENSE)
42
42
  [![Code Style: Ruff](https://img.shields.io/badge/code%20style-ruff-000000.svg)](https://github.com/astral-sh/ruff)
43
+ [![PyPI](https://img.shields.io/pypi/v/codeshield-runtime.svg)](https://pypi.org/project/codeshield-runtime/)
44
+ [![PyPI Downloads](https://img.shields.io/pypi/dm/codeshield-runtime.svg)](https://pypi.org/project/codeshield-runtime/)
43
45
 
44
46
  > Deterministic, Isolated, and Self-Healing Python Execution Runtime for AI Agents.
45
47
 
@@ -142,12 +144,18 @@ pip install -e ".[test,lint,llm,dev]"
142
144
  ### Offline Usage (No API Key)
143
145
 
144
146
  ```python
145
- from codeshield.loop import SelfHealingEngine
147
+ from codeshield import SelfHealingEngine
146
148
 
149
+ # The sandbox is created and destroyed automatically on every ``run`` call.
147
150
  engine = SelfHealingEngine(use_llm=False)
148
- with engine:
149
- result, diagnosis = engine.run("print('hello world')")
150
- print(result.stdout)
151
+ result, diagnosis = engine.run("print('hello world')")
152
+ print(result.stdout)
153
+
154
+ # Use a ``with`` block to reuse a single sandbox across multiple runs.
155
+ with SelfHealingEngine(use_llm=False) as reusable_engine:
156
+ first, _ = reusable_engine.run("print(1 + 1)")
157
+ second, _ = reusable_engine.run("print(2 + 2)")
158
+ print(first.stdout, second.stdout)
151
159
  ```
152
160
 
153
161
  ### Model-Agnostic Self-Healing (Plug-and-Play)
@@ -155,13 +163,13 @@ with engine:
155
163
  CodeShield is not locked into a single LLM. Pass any Python callable as the `patch_generator` to use OpenAI, Anthropic Claude, DeepSeek, Ollama, LiteLLM or your own service:
156
164
 
157
165
  ```python
158
- from codeshield.loop import SelfHealingEngine
166
+ from codeshield import SelfHealingEngine
159
167
 
160
168
 
161
- def custom_openai_patcher(code: str, diagnosis) -> str:
162
- # Any LLM call (OpenAI, Anthropic, DeepSeek, Ollama, LiteLLM)
169
+ def custom_llm_patcher(code: str, diagnosis) -> str:
170
+ # Compatible with any frontier provider: GPT-5.6, Claude Sonnet 5, DeepSeek V4, Ollama
163
171
  response = client.chat.completions.create(
164
- model="gpt-4o-mini",
172
+ model="gpt-5.6-luna", # or "claude-sonnet-5", "deepseek-v4-flash"
165
173
  messages=[
166
174
  {
167
175
  "role": "user",
@@ -172,7 +180,12 @@ def custom_openai_patcher(code: str, diagnosis) -> str:
172
180
  return response.choices[0].message.content
173
181
 
174
182
 
175
- engine = SelfHealingEngine(patch_generator=custom_openai_patcher)
183
+ engine = SelfHealingEngine(patch_generator=custom_llm_patcher)
184
+
185
+ # Broken code -> AST gate -> sandbox execution -> traceback classification ->
186
+ # custom patch -> AST re-validation -> re-execution, all in a single call.
187
+ result, diagnosis = engine.run('print("Result: " + 42)')
188
+ print(result.stdout) # Result: 42
176
189
  ```
177
190
 
178
191
  ### Zero-Config Self-Healing with Gemini Flash
@@ -186,7 +199,8 @@ GEMINI_MODEL=gemini-3.7-flash
186
199
 
187
200
  ```python
188
201
  from dotenv import load_dotenv
189
- from codeshield.loop import SelfHealingEngine
202
+
203
+ from codeshield import SelfHealingEngine
190
204
 
191
205
  load_dotenv()
192
206
 
@@ -218,12 +232,41 @@ python -m codeshield run script.py --no-llm # force local fallback
218
232
  ```python
219
233
  from codeshield import create_code_execution_tool
220
234
 
221
- # Pass the tool directly to your agent
235
+ # Standard usage: built-in Gemini healing when configured, local heuristic otherwise
222
236
  tools = [create_code_execution_tool()]
237
+
238
+ # Or bring your own model: the patcher is forwarded to the internal engine
239
+ tools = [create_code_execution_tool(patch_generator=custom_llm_patcher)]
223
240
  ```
224
241
 
225
242
  `create_code_execution_tool()` returns a ready-to-register `execute_python_code(code: str) -> str` function. It runs the provided Python in a self-healing sandbox and returns either the stdout or a structured error report with `error_type` and `stderr`.
226
243
 
244
+ The callable exposes real type hints and a Google-style docstring, so any SDK that builds a function schema from a plain Python callable can register it directly.
245
+
246
+ ---
247
+
248
+ ## Public API
249
+
250
+ Everything is re-exported at the package root, so imports never need internal submodules:
251
+
252
+ ```python
253
+ from codeshield import (
254
+ ASTSecurityError,
255
+ CodeExecutionRequest,
256
+ ErrorDiagnosis,
257
+ ExecutionResult,
258
+ SandboxManager,
259
+ SelfHealingEngine,
260
+ SelfHealingError,
261
+ SubprocessRunner,
262
+ TracebackClassifier,
263
+ create_code_execution_tool,
264
+ validate_syntax_and_safety,
265
+ )
266
+ ```
267
+
268
+ `ASTSecurityError` subclasses `SelfHealingError` and is raised whenever the static AST gate blocks either the original snippet or a generated patch.
269
+
227
270
  ---
228
271
 
229
272
  ## Verified Examples
@@ -233,18 +276,20 @@ The `examples/` folder contains ready-to-run recipes that have been executed and
233
276
  - `01_basic_sandboxing.py`: isolated execution with timing measurements.
234
277
  - `02_security_gatekeeper.py`: AST rejection of unsafe code.
235
278
  - `03_llm_healing_workflow.py`: self-healing workflow with an LLM or local fallback.
236
- - `04_agent_tool_dropin.py`: end-to-end agentic tool-calling workflow with dynamic code generation.
279
+ - `04_agent_tool_dropin.py`: dual-phase agentic trace, printing agent thought, tool call, sandbox runtime and final answer for both a legitimate analytics round and a security-defense round where the AST gate blocks a shell escape.
280
+ - `05_custom_llm_openai_compatible.py`: model-agnostic, API-key-free self-healing with a custom `patch_generator`.
237
281
 
238
282
  ```bash
239
283
  python examples/01_basic_sandboxing.py
240
284
  python examples/02_security_gatekeeper.py
241
285
  python examples/03_llm_healing_workflow.py
242
286
  python examples/04_agent_tool_dropin.py
287
+ python examples/05_custom_llm_openai_compatible.py
243
288
  ```
244
289
 
245
290
  ## Running Tests & Lint
246
291
 
247
- The suite currently has **50 tests** with **>82% code coverage** on `src/codeshield`.
292
+ The suite currently has **54 tests** with **>83% code coverage** on `src/codeshield`.
248
293
 
249
294
  ```bash
250
295
  ruff check src tests examples
@@ -4,6 +4,8 @@
4
4
  [![Python](https://img.shields.io/badge/python-3.11%20%7C%203.12-blue)](https://www.python.org/)
5
5
  [![License: MIT](https://img.shields.io/badge/license-MIT-green)](LICENSE)
6
6
  [![Code Style: Ruff](https://img.shields.io/badge/code%20style-ruff-000000.svg)](https://github.com/astral-sh/ruff)
7
+ [![PyPI](https://img.shields.io/pypi/v/codeshield-runtime.svg)](https://pypi.org/project/codeshield-runtime/)
8
+ [![PyPI Downloads](https://img.shields.io/pypi/dm/codeshield-runtime.svg)](https://pypi.org/project/codeshield-runtime/)
7
9
 
8
10
  > Deterministic, Isolated, and Self-Healing Python Execution Runtime for AI Agents.
9
11
 
@@ -106,12 +108,18 @@ pip install -e ".[test,lint,llm,dev]"
106
108
  ### Offline Usage (No API Key)
107
109
 
108
110
  ```python
109
- from codeshield.loop import SelfHealingEngine
111
+ from codeshield import SelfHealingEngine
110
112
 
113
+ # The sandbox is created and destroyed automatically on every ``run`` call.
111
114
  engine = SelfHealingEngine(use_llm=False)
112
- with engine:
113
- result, diagnosis = engine.run("print('hello world')")
114
- print(result.stdout)
115
+ result, diagnosis = engine.run("print('hello world')")
116
+ print(result.stdout)
117
+
118
+ # Use a ``with`` block to reuse a single sandbox across multiple runs.
119
+ with SelfHealingEngine(use_llm=False) as reusable_engine:
120
+ first, _ = reusable_engine.run("print(1 + 1)")
121
+ second, _ = reusable_engine.run("print(2 + 2)")
122
+ print(first.stdout, second.stdout)
115
123
  ```
116
124
 
117
125
  ### Model-Agnostic Self-Healing (Plug-and-Play)
@@ -119,13 +127,13 @@ with engine:
119
127
  CodeShield is not locked into a single LLM. Pass any Python callable as the `patch_generator` to use OpenAI, Anthropic Claude, DeepSeek, Ollama, LiteLLM or your own service:
120
128
 
121
129
  ```python
122
- from codeshield.loop import SelfHealingEngine
130
+ from codeshield import SelfHealingEngine
123
131
 
124
132
 
125
- def custom_openai_patcher(code: str, diagnosis) -> str:
126
- # Any LLM call (OpenAI, Anthropic, DeepSeek, Ollama, LiteLLM)
133
+ def custom_llm_patcher(code: str, diagnosis) -> str:
134
+ # Compatible with any frontier provider: GPT-5.6, Claude Sonnet 5, DeepSeek V4, Ollama
127
135
  response = client.chat.completions.create(
128
- model="gpt-4o-mini",
136
+ model="gpt-5.6-luna", # or "claude-sonnet-5", "deepseek-v4-flash"
129
137
  messages=[
130
138
  {
131
139
  "role": "user",
@@ -136,7 +144,12 @@ def custom_openai_patcher(code: str, diagnosis) -> str:
136
144
  return response.choices[0].message.content
137
145
 
138
146
 
139
- engine = SelfHealingEngine(patch_generator=custom_openai_patcher)
147
+ engine = SelfHealingEngine(patch_generator=custom_llm_patcher)
148
+
149
+ # Broken code -> AST gate -> sandbox execution -> traceback classification ->
150
+ # custom patch -> AST re-validation -> re-execution, all in a single call.
151
+ result, diagnosis = engine.run('print("Result: " + 42)')
152
+ print(result.stdout) # Result: 42
140
153
  ```
141
154
 
142
155
  ### Zero-Config Self-Healing with Gemini Flash
@@ -150,7 +163,8 @@ GEMINI_MODEL=gemini-3.7-flash
150
163
 
151
164
  ```python
152
165
  from dotenv import load_dotenv
153
- from codeshield.loop import SelfHealingEngine
166
+
167
+ from codeshield import SelfHealingEngine
154
168
 
155
169
  load_dotenv()
156
170
 
@@ -182,12 +196,41 @@ python -m codeshield run script.py --no-llm # force local fallback
182
196
  ```python
183
197
  from codeshield import create_code_execution_tool
184
198
 
185
- # Pass the tool directly to your agent
199
+ # Standard usage: built-in Gemini healing when configured, local heuristic otherwise
186
200
  tools = [create_code_execution_tool()]
201
+
202
+ # Or bring your own model: the patcher is forwarded to the internal engine
203
+ tools = [create_code_execution_tool(patch_generator=custom_llm_patcher)]
187
204
  ```
188
205
 
189
206
  `create_code_execution_tool()` returns a ready-to-register `execute_python_code(code: str) -> str` function. It runs the provided Python in a self-healing sandbox and returns either the stdout or a structured error report with `error_type` and `stderr`.
190
207
 
208
+ The callable exposes real type hints and a Google-style docstring, so any SDK that builds a function schema from a plain Python callable can register it directly.
209
+
210
+ ---
211
+
212
+ ## Public API
213
+
214
+ Everything is re-exported at the package root, so imports never need internal submodules:
215
+
216
+ ```python
217
+ from codeshield import (
218
+ ASTSecurityError,
219
+ CodeExecutionRequest,
220
+ ErrorDiagnosis,
221
+ ExecutionResult,
222
+ SandboxManager,
223
+ SelfHealingEngine,
224
+ SelfHealingError,
225
+ SubprocessRunner,
226
+ TracebackClassifier,
227
+ create_code_execution_tool,
228
+ validate_syntax_and_safety,
229
+ )
230
+ ```
231
+
232
+ `ASTSecurityError` subclasses `SelfHealingError` and is raised whenever the static AST gate blocks either the original snippet or a generated patch.
233
+
191
234
  ---
192
235
 
193
236
  ## Verified Examples
@@ -197,18 +240,20 @@ The `examples/` folder contains ready-to-run recipes that have been executed and
197
240
  - `01_basic_sandboxing.py`: isolated execution with timing measurements.
198
241
  - `02_security_gatekeeper.py`: AST rejection of unsafe code.
199
242
  - `03_llm_healing_workflow.py`: self-healing workflow with an LLM or local fallback.
200
- - `04_agent_tool_dropin.py`: end-to-end agentic tool-calling workflow with dynamic code generation.
243
+ - `04_agent_tool_dropin.py`: dual-phase agentic trace, printing agent thought, tool call, sandbox runtime and final answer for both a legitimate analytics round and a security-defense round where the AST gate blocks a shell escape.
244
+ - `05_custom_llm_openai_compatible.py`: model-agnostic, API-key-free self-healing with a custom `patch_generator`.
201
245
 
202
246
  ```bash
203
247
  python examples/01_basic_sandboxing.py
204
248
  python examples/02_security_gatekeeper.py
205
249
  python examples/03_llm_healing_workflow.py
206
250
  python examples/04_agent_tool_dropin.py
251
+ python examples/05_custom_llm_openai_compatible.py
207
252
  ```
208
253
 
209
254
  ## Running Tests & Lint
210
255
 
211
- The suite currently has **50 tests** with **>82% code coverage** on `src/codeshield`.
256
+ The suite currently has **54 tests** with **>83% code coverage** on `src/codeshield`.
212
257
 
213
258
  ```bash
214
259
  ruff check src tests examples
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "codeshield-runtime"
7
- version = "0.1.1"
7
+ version = "0.1.3"
8
8
  description = "A secure, isolated, and self-healing Python code execution engine powered by uv and AST analysis."
9
9
  readme = "README.md"
10
10
  license = { text = "MIT" }
@@ -0,0 +1,46 @@
1
+ """CodeShield.
2
+
3
+ A secure, isolated, and self-healing Python code execution engine.
4
+ """
5
+
6
+ __version__ = "0.1.3"
7
+
8
+ from codeshield.analyzer import validate_syntax_and_safety
9
+ from codeshield.classifier import TracebackClassifier
10
+ from codeshield.environment import SandboxError, SandboxManager
11
+ from codeshield.loop import (
12
+ ASTSecurityError,
13
+ GeminiPatchGenerator,
14
+ LLMPatchError,
15
+ SelfHealingEngine,
16
+ SelfHealingError,
17
+ )
18
+ from codeshield.runner import SubprocessRunner, SubprocessRunnerError
19
+ from codeshield.schemas import (
20
+ CodeExecutionRequest,
21
+ ErrorDiagnosis,
22
+ ExecutionResult,
23
+ PatchProposal,
24
+ ValidationReport,
25
+ )
26
+ from codeshield.tools import create_code_execution_tool
27
+
28
+ __all__ = [
29
+ "ASTSecurityError",
30
+ "CodeExecutionRequest",
31
+ "ErrorDiagnosis",
32
+ "ExecutionResult",
33
+ "GeminiPatchGenerator",
34
+ "LLMPatchError",
35
+ "PatchProposal",
36
+ "SandboxError",
37
+ "SandboxManager",
38
+ "SelfHealingEngine",
39
+ "SelfHealingError",
40
+ "SubprocessRunner",
41
+ "SubprocessRunnerError",
42
+ "TracebackClassifier",
43
+ "ValidationReport",
44
+ "create_code_execution_tool",
45
+ "validate_syntax_and_safety",
46
+ ]
@@ -65,7 +65,7 @@ class SandboxManager:
65
65
 
66
66
  @property
67
67
  def workspace_path(self) -> Path:
68
- """Alias for the root sandbox workspace directory."""
68
+ """Alias of :attr:`workspace` used for explicit lifecycle checks."""
69
69
  return self._workspace
70
70
 
71
71
  @property
@@ -200,6 +200,10 @@ class SandboxManager:
200
200
  ) -> Path:
201
201
  """Write package specifiers to ``workspace/requirements.txt``.
202
202
 
203
+ Args:
204
+ packages: Package specifiers to write, one per line.
205
+ file_name: Name of the requirements file inside the workspace.
206
+
203
207
  Returns:
204
208
  Path of the written file.
205
209
  """
@@ -218,10 +222,12 @@ class SandboxManager:
218
222
  logger.warning("Could not remove workspace %s: %s", self._workspace, exc)
219
223
 
220
224
  def __enter__(self) -> SandboxManager:
225
+ """Create the virtual environment and return the manager."""
221
226
  self.create()
222
227
  return self
223
228
 
224
229
  def __exit__(self, *exc: object) -> None:
230
+ """Remove the workspace when leaving the context."""
225
231
  self.cleanup()
226
232
 
227
233
  def _resolve_backend(self, backend: str | None) -> str:
@@ -1,3 +1,4 @@
1
+ # ruff: noqa: UP045
1
2
  """Self-healing execution loop."""
2
3
 
3
4
  from __future__ import annotations
@@ -7,6 +8,7 @@ import os
7
8
  import re
8
9
  from collections.abc import Callable
9
10
  from pathlib import Path
11
+ from typing import Optional
10
12
 
11
13
  from codeshield.analyzer import validate_syntax_and_safety
12
14
  from codeshield.classifier import TracebackClassifier
@@ -32,6 +34,10 @@ class SelfHealingError(RuntimeError):
32
34
  """Raised when the self-healing loop cannot complete execution safely."""
33
35
 
34
36
 
37
+ class ASTSecurityError(SelfHealingError):
38
+ """Raised when code or a generated patch is blocked by the static AST gate."""
39
+
40
+
35
41
  class LLMPatchError(RuntimeError):
36
42
  """Raised when the LLM patch generator cannot produce a safe correction."""
37
43
 
@@ -173,7 +179,7 @@ class SelfHealingEngine:
173
179
  sandbox: SandboxManager | None = None,
174
180
  runner: SubprocessRunner | None = None,
175
181
  classifier: TracebackClassifier | None = None,
176
- patch_generator: Callable[[str, ErrorDiagnosis], str | None] | None = None,
182
+ patch_generator: Optional[Callable[[str, ErrorDiagnosis], Optional[str]]] = None,
177
183
  gemini_api_key: str | None = None,
178
184
  gemini_model: str | None = None,
179
185
  use_llm: bool = True,
@@ -237,6 +243,8 @@ class SelfHealingEngine:
237
243
  Raises:
238
244
  SelfHealingError: when the loop exhausts all iterations without a
239
245
  clean result.
246
+ ASTSecurityError: when the code or a generated patch is blocked by
247
+ the static AST gate.
240
248
  """
241
249
  if isinstance(request, str):
242
250
  request = CodeExecutionRequest(code=request)
@@ -274,7 +282,7 @@ class SelfHealingEngine:
274
282
  )
275
283
  patched = self._generate_and_validate_patch(request.code, diagnosis)
276
284
  if patched is None:
277
- raise SelfHealingError(
285
+ raise ASTSecurityError(
278
286
  f"Static safety violations cannot be auto-patched: {report.violations}"
279
287
  )
280
288
  request = self._apply_patch(request, patched, attempt)
@@ -335,9 +343,21 @@ class SelfHealingEngine:
335
343
  proposal: PatchProposal,
336
344
  attempt: int,
337
345
  ) -> CodeExecutionRequest:
338
- """Return a new request with the patched code and an updated file name."""
346
+ """Return a new request with the patched code and an updated file name.
347
+
348
+ Args:
349
+ request: The request that produced the failing execution.
350
+ proposal: The AST-validated patch proposal to apply.
351
+ attempt: Current healing iteration, used to name the patched script.
352
+
353
+ Returns:
354
+ A new ``CodeExecutionRequest`` carrying the patched source.
355
+
356
+ Raises:
357
+ ASTSecurityError: when the proposal failed static AST validation.
358
+ """
339
359
  if not proposal.is_syntax_valid:
340
- raise SelfHealingError(
360
+ raise ASTSecurityError(
341
361
  f"Proposed patch failed AST validation: {proposal.patched_code[:200]}"
342
362
  )
343
363
 
@@ -10,6 +10,7 @@ import subprocess
10
10
  import threading
11
11
  import time
12
12
  from collections.abc import Iterable
13
+ from typing import IO
13
14
 
14
15
  from codeshield.environment import SandboxError, SandboxManager
15
16
  from codeshield.schemas import CodeExecutionRequest, ExecutionResult
@@ -146,7 +147,7 @@ class SubprocessRunner:
146
147
  stdout_lines: list[str] = []
147
148
  stderr_lines: list[str] = []
148
149
 
149
- def reader(pipe, sink: list[str]) -> None:
150
+ def reader(pipe: IO[str], sink: list[str]) -> None:
150
151
  """Read lines from a pipe without blocking the main thread."""
151
152
  try:
152
153
  for line in pipe:
@@ -1,28 +1,36 @@
1
+ # ruff: noqa: UP045
1
2
  """Universal agent tool wrapper for CodeShield.
2
3
 
3
4
  The function returned by ``create_code_execution_tool`` can be registered as a
4
5
  tool in any agent framework (LangChain, CrewAI, Google Gen AI, etc.). It runs
5
6
  the provided Python source inside a self-healing sandbox and returns either the
6
7
  stdout or a structured error report.
7
- """
8
8
 
9
- from __future__ import annotations
9
+ Annotations in this module are intentionally *not* postponed: agent SDKs such as
10
+ ``google-genai`` introspect ``execute_python_code.__annotations__`` at runtime to
11
+ build the function-calling schema, and string annotations break that conversion.
12
+ """
10
13
 
11
14
  from collections.abc import Callable
15
+ from typing import Optional
12
16
 
13
17
  from codeshield.environment import SandboxError
14
18
  from codeshield.loop import SelfHealingEngine, SelfHealingError
15
19
  from codeshield.runner import SubprocessRunnerError
20
+ from codeshield.schemas import ErrorDiagnosis
16
21
 
17
22
 
18
23
  def create_code_execution_tool(
19
- engine: SelfHealingEngine | None = None,
24
+ engine: Optional[SelfHealingEngine] = None,
25
+ patch_generator: Optional[Callable[[str, ErrorDiagnosis], Optional[str]]] = None,
20
26
  ) -> Callable[[str], str]:
21
27
  """Return a drop-in ``execute_python_code(code: str) -> str`` tool.
22
28
 
23
29
  Args:
24
30
  engine: Optional ``SelfHealingEngine`` instance. When ``None``, a fresh
25
31
  engine is created for each tool call.
32
+ patch_generator: Optional custom patcher ``(code, diagnosis) -> patched``
33
+ forwarded to ``SelfHealingEngine`` when ``engine`` is not provided.
26
34
 
27
35
  Returns:
28
36
  A callable ready to be registered as an agent tool.
@@ -41,13 +49,11 @@ def create_code_execution_tool(
41
49
  The stdout of the script if execution succeeds, or a structured
42
50
  error report if it fails after all self-healing attempts.
43
51
  """
44
- _engine = engine or SelfHealingEngine()
52
+ _engine = engine or SelfHealingEngine(patch_generator=patch_generator)
45
53
  with _engine:
46
54
  try:
47
55
  result, diagnosis = _engine.run(code)
48
- except SelfHealingError as exc:
49
- return f"error_type: SelfHealingError\nmessage: {exc}"
50
- except (SandboxError, SubprocessRunnerError) as exc:
56
+ except (SelfHealingError, SandboxError, SubprocessRunnerError) as exc:
51
57
  return f"error_type: {type(exc).__name__}\nmessage: {exc}"
52
58
 
53
59
  if (
@@ -9,15 +9,17 @@ from unittest.mock import MagicMock
9
9
 
10
10
  import pytest
11
11
 
12
- from codeshield.classifier import TracebackClassifier
13
- from codeshield.environment import SandboxError, SandboxManager
14
- from codeshield.loop import (
12
+ from codeshield import (
13
+ CodeExecutionRequest,
14
+ ErrorDiagnosis,
15
15
  GeminiPatchGenerator,
16
+ SandboxError,
17
+ SandboxManager,
16
18
  SelfHealingEngine,
17
19
  SelfHealingError,
20
+ SubprocessRunner,
21
+ TracebackClassifier,
18
22
  )
19
- from codeshield.runner import SubprocessRunner
20
- from codeshield.schemas import CodeExecutionRequest, ErrorDiagnosis
21
23
 
22
24
 
23
25
  @pytest.fixture
@@ -4,13 +4,14 @@ from __future__ import annotations
4
4
 
5
5
  import pytest
6
6
 
7
- from codeshield.analyzer import validate_syntax_and_safety
8
- from codeshield.classifier import TracebackClassifier
9
- from codeshield.environment import SandboxManager
10
- from codeshield.loop import SelfHealingEngine
11
- from codeshield.runner import SubprocessRunner
12
- from codeshield.schemas import (
7
+ from codeshield import (
8
+ ASTSecurityError,
13
9
  ExecutionResult,
10
+ SandboxManager,
11
+ SelfHealingEngine,
12
+ SubprocessRunner,
13
+ TracebackClassifier,
14
+ validate_syntax_and_safety,
14
15
  )
15
16
 
16
17
 
@@ -244,3 +245,38 @@ def test_self_healing_engine_run_without_context_manager() -> None:
244
245
  assert result.stdout.strip() == "no context manager"
245
246
  assert diagnosis is None
246
247
  assert not engine._sandbox.workspace_path.exists()
248
+
249
+
250
+ def test_custom_patch_generator_success() -> None:
251
+ """A custom patch_generator can repair code that the local fallback cannot."""
252
+
253
+ def custom_patcher(code: str, diagnosis) -> str | None:
254
+ if diagnosis.error_type == "TypeError":
255
+ return code.replace('"Result: " + 42', '"Result: " + str(42)')
256
+ return None
257
+
258
+ engine = SelfHealingEngine(
259
+ patch_generator=custom_patcher,
260
+ use_llm=False,
261
+ )
262
+ with engine:
263
+ result, diagnosis = engine.run('print("Result: " + 42)')
264
+
265
+ assert result.exit_code == 0
266
+ assert result.stdout.strip() == "Result: 42"
267
+ assert diagnosis is None
268
+
269
+
270
+ def test_custom_patch_generator_ast_rejection() -> None:
271
+ """A custom patch that violates the AST security gate is rejected."""
272
+
273
+ def malicious_patcher(code: str, diagnosis) -> str | None:
274
+ return 'import os\nos.system("echo pwned")'
275
+
276
+ engine = SelfHealingEngine(
277
+ patch_generator=malicious_patcher,
278
+ use_llm=False,
279
+ )
280
+
281
+ with pytest.raises(ASTSecurityError):
282
+ engine.run("print(undefined_value)")
@@ -2,8 +2,16 @@
2
2
 
3
3
  from __future__ import annotations
4
4
 
5
- from codeshield import create_code_execution_tool
6
- from codeshield.loop import SelfHealingEngine
5
+ from codeshield import SelfHealingEngine, create_code_execution_tool
6
+
7
+
8
+ def test_tool_exposes_runtime_type_annotations() -> None:
9
+ """Agent SDKs build the schema from ``__annotations__``, so they must be types."""
10
+ tool = create_code_execution_tool()
11
+
12
+ assert tool.__name__ == "execute_python_code"
13
+ assert tool.__annotations__["code"] is str
14
+ assert tool.__annotations__["return"] is str
7
15
 
8
16
 
9
17
  def test_tool_returns_stdout_on_success() -> None:
@@ -1,17 +0,0 @@
1
- """CodeShield.
2
-
3
- A secure, isolated, and self-healing Python code execution engine.
4
- """
5
-
6
- __version__ = "0.1.0"
7
-
8
- from codeshield.loop import SelfHealingEngine
9
- from codeshield.schemas import CodeExecutionRequest, ExecutionResult
10
- from codeshield.tools import create_code_execution_tool
11
-
12
- __all__ = [
13
- "SelfHealingEngine",
14
- "CodeExecutionRequest",
15
- "ExecutionResult",
16
- "create_code_execution_tool",
17
- ]