refactor: modernize engine with Pydantic models, modular architecture, and session recovery#166
Open
anticomputer wants to merge 19 commits intomainfrom
Open
refactor: modernize engine with Pydantic models, modular architecture, and session recovery#166anticomputer wants to merge 19 commits intomainfrom
anticomputer wants to merge 19 commits intomainfrom
Conversation
anticomputer
requested review from
JarLob,
kevinbackhouse,
m-y-mo and
p-
as code owners
anticomputer
requested review from
Kwstubbs and
sylwia-budzynska
as code owners
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Pull request overview
Refactors the SecLab Taskflow Agent from a monolithic entrypoint into a modular engine with Pydantic v2-validated YAML grammar models, a Typer CLI, per-model API configuration (Chat Completions vs Responses), and task-level checkpoint/resume support.
Changes:
- Introduces typed Pydantic document models and updates YAML loading to validate at parse time.
- Replaces the old argparse-based entrypoint with a Typer CLI and moves execution logic into a dedicated runner with MCP lifecycle modules.
- Adds session checkpointing/resume plus new examples/docs/tests for the modernized architecture and per-model API settings.
Reviewed changes
Copilot reviewed 35 out of 35 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/test_yaml_parser.py | Updates YAML parsing assertions to use Pydantic document attributes. |
| tests/test_template_utils.py | Switches reusable taskflow prompt access to typed models. |
| tests/test_session.py | Adds unit tests for session checkpoint/load/list behavior. |
| tests/test_models.py | Adds validation tests for Pydantic grammar models and real YAML fixtures. |
| tests/test_cli_parser.py | Updates legacy prompt parsing tests for renamed unpack variables and typed taskflow access. |
| src/seclab_taskflow_agent/template_utils.py | Adds typing/exports and adapts prompt loading to Pydantic PromptDocument. |
| src/seclab_taskflow_agent/shell_utils.py | Adds typing/docs and improves shell error messaging. |
| src/seclab_taskflow_agent/session.py | Implements taskflow session persistence for checkpoint/resume. |
| src/seclab_taskflow_agent/runner.py | New modular execution engine with model resolution, reusable tasks, prompt expansion, retries, and session handling. |
| src/seclab_taskflow_agent/render_utils.py | Adds typing/exports and docs for async output buffering. |
| src/seclab_taskflow_agent/prompt_parser.py | Extracts legacy argparse parsing used by runtime/test compatibility. |
| src/seclab_taskflow_agent/path_utils.py | Centralizes platformdirs data-dir creation and reuses it for MCP data storage. |
| src/seclab_taskflow_agent/models.py | Defines the YAML grammar via Pydantic v2 models (taskflow/personality/toolbox/model_config/prompt). |
| src/seclab_taskflow_agent/mcp_utils.py | Slims MCP utilities and re-exports transport/prompt; updates toolbox param resolution for typed models. |
| src/seclab_taskflow_agent/mcp_transport.py | Extracts MCP transport implementations (stdio wrappers + streamable process thread). |
| src/seclab_taskflow_agent/mcp_servers/memcache/memcache_backend/dictionary_file.py | Tightens type comparisons for memcache “add_state”. |
| src/seclab_taskflow_agent/mcp_servers/codeql/jsonrpyc/init.py | Replaces bare except with except Exception. |
| src/seclab_taskflow_agent/mcp_servers/codeql/client.py | Adds logging import (used for client logging). |
| src/seclab_taskflow_agent/mcp_prompt.py | Extracts system prompt construction helper. |
| src/seclab_taskflow_agent/mcp_lifecycle.py | Extracts MCP server connect/cleanup lifecycle management. |
| src/seclab_taskflow_agent/env_utils.py | Adds typing/docs and exports for env swap + temporary env context manager. |
| src/seclab_taskflow_agent/cli.py | Adds Typer CLI with --debug/--resume, concise error printing, and logging setup. |
| src/seclab_taskflow_agent/capi.py | Adds module docs/exports and typing improvements for token/endpoint/model listing. |
| src/seclab_taskflow_agent/banner.py | Adds typing/exports and docstring for startup banner. |
| src/seclab_taskflow_agent/available_tools.py | Rewrites YAML loader to cache + validate documents via Pydantic models. |
| src/seclab_taskflow_agent/agent.py | Extends TaskAgent to support Responses API, per-model endpoint/token overrides, and adds exports/docs. |
| src/seclab_taskflow_agent/main.py | Replaces monolithic implementation with a small entrypoint delegating to CLI/runner. |
| src/seclab_taskflow_agent/init.py | Adds package-level docs and re-exports key public types. |
| pyproject.toml | Updates metadata, adds console script entrypoint, and tightens Ruff configuration. |
| examples/taskflows/example_responses_api.yaml | Adds an example taskflow configured to use the Responses API. |
| examples/taskflows/echo_responses_api.yaml | Adds a Responses API echo taskflow example. |
| examples/taskflows/comprehensive_test.yaml | Adds a comprehensive taskflow exercising the full grammar and engine features. |
| examples/model_configs/responses_api.yaml | Adds model_config demonstrating per-model api_type/endpoint/token overrides. |
| doc/GRAMMAR.md | Documents per-model settings and engine-handled model_settings keys. |
| README.md | Documents the new architecture, API types, session recovery, and error output behavior. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
You can also share your feedback on Copilot code review. Take the survey.
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 35 out of 35 changed files in this pull request and generated 4 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
You can also share your feedback on Copilot code review. Take the survey.
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 36 out of 36 changed files in this pull request and generated 4 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
You can also share your feedback on Copilot code review. Take the survey.
p-
reviewed
Mar 16, 2026
…r CLI - Extract __main__.py into cli.py, runner.py, mcp_lifecycle.py, models.py - Add Pydantic v2 grammar models for all YAML document types - Replace argparse with Typer CLI, add project.scripts entry point - Reduce ruff ignore list from 59 to 22 rules, fix lint issues - Add 29 model tests against real YAML files (68/68 pass) - Full grammar backwards compatibility preserved
Replace all raw dict access with typed model attributes throughout the engine. AvailableTools now returns Pydantic model instances, runner.py uses typed fields, and tests validate via model attributes.
Split mcp_utils into mcp_transport and mcp_prompt. Extract prompt_parser to break cli/runner circular import. Decompose run_main into focused helpers. Add type hints and docstrings across all modules. Add __all__ exports, pyproject keywords and classifiers.
Add api_type field to ModelConfigDocument (chat_completions|responses). Thread api_type through runner -> deploy_task_agents -> TaskAgent. TaskAgent switches between OpenAIChatCompletionsModel and OpenAIResponsesModel based on api_type. Default: chat_completions.
model_settings in model_config can now set api_type, endpoint, and token per model. token is an env var name to resolve. Allows mixing chat_completions and responses API across models and routing to different endpoints within a single taskflow.
Document api_type, endpoint, and token per-model overrides. Update architecture module listing in README. Add per-model settings reference table to GRAMMAR.md.
- Fix mutable default args in TaskAgent.__init__ (list -> None) - Replace bare except Exception with specific exception types - Add __all__ exports to all 14 submodules - Add docstrings to all hook methods in agent.py - Add return type to banner.py and prompt_parser.py - Add module docstrings to capi.py, env_utils.py, shell_utils.py - Remove B006 ruff suppression (no longer needed)
Add --debug/-d flag and TASK_AGENT_DEBUG env var. Default: one-line error chain. Debug: full traceback.
Add TaskflowSession model for task-level checkpointing. Save progress after each task, auto-retry failed tasks 3x. New --resume flag to continue from last checkpoint. Includes 9 tests for session persistence.
Covers: shell tasks, repeat_prompt, async iteration, model_config, globals, inputs, env, MCP toolboxes, headless, blocked_tools, reusable tasks (uses), reusable prompts (include), and handoffs.
- fix uninitialized args and inconsistent return shape in prompt_parser - add explicit returns after exhaustive match statements in capi - narrow BaseException to Exception in mcp_transport thread - add comment for empty except in cleanup shutdown path - iterate over copy in mcp_lifecycle cleanup to avoid mutation during iteration - remove unused _ENGINE_SETTING_KEYS constant from runner - remove unused import in test_session - default TaskDefinition.name/description to empty string for per-index naming - use explicit re-export pattern (as X) in __main__ and cli
- remove global OpenAI client/api mutations (race condition with concurrent agents) - add empty resolved_agents validation before deploy_task_agents - store toolbox name on MCPServerEntry instead of accessing private _name - update last_tool_results unconditionally in session.record_task - move raises into match default cases to fix mixed return warnings - catch non-SystemExit exceptions in prompt parser
- don't advance resume cursor past failed must_complete tasks - peek at last tool result before consuming (safe for retry) - use Traversable.open() for zip/wheel compatibility - return empty string (not None) for prompt on invalid globals
- cli: TASK_AGENT_DEBUG="0"/"false" no longer enables debug mode - capi: allow arbitrary API endpoints with graceful fallback - runner: defer tool result pop until after template rendering - test: 72 new unit tests for runner, cli, session, prompt parser, capi - examples: add edge_case_test.yaml for nested JSON repeat_prompt
anticomputer
force-pushed
the
anticomputer/refactor
branch
from
c22b38f to
69fc1e3
Compare
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 43 out of 43 changed files in this pull request and generated 5 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
- capi: unknown endpoint fallback now type-checks response before iterating, avoids AttributeError on dict responses without data key - render_utils: flush_async_output is a no-op when no buffer exists, prevents masking the real error when an async agent fails early - session: list_sessions sorts by file mtime instead of filename so most-recent-first ordering is actually correct - runner: replace confusing issubset(set()) idiom with straightforward set difference check for unknown model settings keys - runner: catch JSONDecodeError on outer tool result parse separately from the inner text parse so malformed results get a clear error - tests: suppress S105 false positive on test token assertion
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 43 out of 43 changed files in this pull request and generated 2 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The engine code was a single 671-line
__main__.pywith raw dict access everywhere, no type safety, and no error recovery. This rewrites the internals while keeping the YAML grammar fully backwards-compatible. All existing taskflows (both example and production seclab-taskflows) work without changes.What changed
Architecture.
__main__.pyis now 25 lines. The logic lives in focused modules:models.py-- Pydantic v2 models for all five YAML document types (taskflow, personality, toolbox, model_config, prompt). Validation happens at parse time instead of scattering.get()calls everywhere.runner.py-- execution engine, extracted from__main__and decomposed into_resolve_model_config(),_merge_reusable_task(),_resolve_task_model(),_build_prompts_to_run().cli.py-- Typer-based CLI replacing argparse. Adds--debugand--resumeflags.agent.py-- TaskAgent wrapper, now supports per-model API type, endpoint, and token overrides.mcp_lifecycle.py,mcp_transport.py,mcp_prompt.py-- decomposed from the old 462-linemcp_utils.py.session.py-- task-level checkpoint/resume with auto-retry.prompt_parser.py-- extracted to break a circular import between cli and runner.Pydantic grammar models. Every YAML file is validated against a typed model at load time.
extra="allow"on all models so new fields do not break old parsers. Reserved word conflicts handled via aliases (async->async_task,model_config->model_config_ref). Version normalization (1->"1.0") with validation.Per-model API configuration.
model_settingsin a model_config can now includeapi_type(chat_completions or responses),endpoint, andtoken(env var name). These are extracted by the engine before passing remaining settings to the SDK. This lets you mix endpoints and API types across models in a single taskflow.Arbitrary API endpoints. Unknown API endpoints no longer crash the agent at import time or when listing models. The model catalog and tool-call detection fall back gracefully for non-standard endpoints, so users can point at any compatible API.
Session recovery. Taskflow runs are checkpointed after each task. On unrecoverable failure, the session ID is printed and
--resumepicks up from the last successful task. Failedmust_completetasks are not advanced past on resume. Failed tasks auto-retry 3 times with backoff before giving up.Error output. Exceptions now print a concise one-line chain by default instead of dumping 80-line tracebacks through httpx/httpcore/openai internals.
--debugorTASK_AGENT_DEBUG=1gets you the full trace. SettingTASK_AGENT_DEBUG=0orTASK_AGENT_DEBUG=falsecorrectly disables debug mode.Code quality. Type hints and docstrings on all public functions.
__all__exports on every module. Fixed mutable default args inTaskAgent.__init__. Replaced bareexcept Exceptionwith specific types. Zero ruff errors across the entire codebase. Ruff ignores reduced from 59 to 21.What did not change
The YAML grammar is 100% backwards-compatible. No taskflow files need modification. All field names, nesting, template syntax, and behavior are preserved.
Testing
Unit tests for Pydantic models, runner internals (model resolution, task merging, prompt building, pop-after-render correctness), Typer CLI arg parsing, debug env handling, prompt parser edge cases, capi endpoint handling with arbitrary endpoint fallback, and session persistence including corrupt file recovery. 155 tests total, all passing, ruff clean.
Integration tested against api.githubcopilot.com: