diff --git a/.claude-plugin/marketplace.json b/.claude-plugin/marketplace.json new file mode 100644 index 00000000..346e42c5 --- /dev/null +++ b/.claude-plugin/marketplace.json @@ -0,0 +1,18 @@ +{ + "name": "flowcept-harness", + "owner": { + "name": "flowcept-harness" + }, + "metadata": { + "description": "PROV-AGENT provenance capture and analysis for AI coding harnesses, backed by Flowcept.", + "version": "0.2.0" + }, + "plugins": [ + { + "name": "flowcept", + "source": "./plugins/flowcept", + "description": "Capture Claude Code sessions as Flowcept PROV-AGENT provenance and analyze them with MCP provenance tools, analysis skills, and optional auto-reports.", + "category": "observability" + } + ] +} diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index b913bcd5..f7739f37 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -19,7 +19,7 @@ jobs: - name: Install package and dependencies run: | pip install --upgrade pip - pip install ruff + pip install ruff==0.15.22 pip install .[docs,webservice,extras] - name: Run linter and formatter checks using ruff @@ -28,6 +28,9 @@ jobs: - name: Run HTML builder for Sphinx documentation run: make docs + - name: Check that the committed OpenAPI spec is up to date + run: git diff --exit-code -- docs/openapi/flowcept-openapi.json docs/openapi/flowcept-openapi.yaml + - name: Clean up run: | make clean diff --git a/.github/workflows/run-tests.yml b/.github/workflows/run-tests.yml index caef619e..757e949e 100644 --- a/.github/workflows/run-tests.yml +++ b/.github/workflows/run-tests.yml @@ -35,7 +35,7 @@ jobs: - name: Install package and dependencies run: | - pip install ruff + pip install ruff==0.15.22 pip install .[docs] - name: List installed packages diff --git a/.gitignore b/.gitignore index c4f94406..385fd558 100644 --- a/.gitignore +++ b/.gitignore @@ -32,7 +32,7 @@ launch.json .vscode/ **/*.err **/*.out -core.* +core.[0-9]* *.csv flowcept_code_assistants_memory.md uv.lock diff --git a/README.md b/README.md index c6f44356..3ccebe35 100644 --- a/README.md +++ b/README.md @@ -202,6 +202,281 @@ Flowcept supports several capture styles. Use the least invasive one that answer Read [Provenance Capture Methods](https://flowcept.readthedocs.io/en/latest/prov_capture.html) for examples. +## Agentic Provenance Plugins + +Flowcept ships with zero-code-change provenance plugins for four popular agentic frameworks: **Academy**, **LangGraph**, **CrewAI**, and **AutoGen**. Each plugin automatically captures: + +- **Intra-agent provenance** — individual action/task executions with inputs, outputs, timing, and status +- **Inter-agent provenance** — parent/child relationships between agents and the tasks they spawn +- **LLM call provenance** — every OpenAI or Anthropic API call linked back to the agent action that triggered it (model, prompt, tokens, latency) + +### Enabling plugins via `settings.yaml` + +Add a `plugins:` block to your `~/.flowcept/settings.yaml`. Only the frameworks you want to track need to be listed: + +```yaml +plugins: + academy: + enabled: true + kind: academy + workflow_name: "my-academy-workflow" + performance_tracking: true + langgraph: + enabled: true + kind: langgraph + workflow_name: "my-langgraph-workflow" + performance_tracking: true + crewai: + enabled: true + kind: crewai + workflow_name: "my-crewai-workflow" + performance_tracking: true + autogen: + enabled: true + kind: autogen + workflow_name: "my-autogen-workflow" + performance_tracking: true +``` + +Then wrap your code with `Flowcept()` — all enabled plugins start and stop automatically: + +```python +from flowcept import Flowcept + +with Flowcept(): + # your Academy / LangGraph / CrewAI / AutoGen code here + ... +``` + +### Helper utilities + +Each plugin exposes drop-in wrappers to record LLM calls regardless of which framework is active: + +| Function / Class | Purpose | +|---|---| +| `openai_chat(prompt, model, ...)` | Call OpenAI and automatically record the call with provenance linkage | +| `anthropic_chat(prompt, model, ...)` | Same for Anthropic Claude models | +| `FlowceptAnthropicClient(client)` | Wrap an existing `anthropic.Anthropic` client to intercept all `messages.create` / `stream` calls | +| `run_team(team, task, ...)` | *(AutoGen only)* Run an autogen team and capture full provenance without an explicit plugin handle | + +All four plugins export `openai_chat` and `anthropic_chat` from their respective modules (e.g. `from flowcept.agents.academy.academy_plugin import openai_chat`). + +### What gets captured + +| Captured field | Academy | LangGraph | CrewAI | AutoGen | +|---|:---:|:---:|:---:|:---:| +| Agent action / node executions | ✓ | ✓ | ✓ | ✓ | +| Inputs and outputs per action | ✓ | ✓ | ✓ | ✓ | +| Timing (start / end / latency) | ✓ | ✓ | ✓ | ✓ | +| Parent–child task linkage | ✓ | ✓ | ✓ | ✓ | +| LLM calls (OpenAI) | ✓ | ✓ | ✓ | ✓ | +| LLM calls (Anthropic) | ✓ | ✓ | ✓ | ✓ | +| Token usage | ✓ | ✓ | ✓ | ✓ | +| Agent ID on LLM calls | ✓ | ✓ | ✓ | ✓ | +| Automatic agent wrapping | ✓ | ✓ | — | ✓ | + +### Provenance record types per plugin + +Each plugin emits typed `TaskObject` records tagged with a `subtype` field, forming a nested provenance hierarchy: + +| Plugin | Record subtypes and hierarchy | +|---|---| +| **Academy** | Campaign `WorkflowObject` → agent `WorkflowObject` → `academy_action` / `academy_loop` / `academy_lifecycle` (siblings) → `llm_call` (under action or loop) | +| **LangGraph** | `WorkflowObject` → `langgraph_graph` → `langgraph_node` → `llm_call` / `tool_call` | +| **CrewAI** | `WorkflowObject` → `crewai_crew` (no children via `parent_task_id`) · `crewai_task` → `crewai_agent` → `llm_call` / `tool_call` | +| **AutoGen** | `WorkflowObject` → `autogen_run` → `autogen_message` → `llm_call` | + +All records carry `campaign_id`, `workflow_id`, `task_id`, `started_at`, `ended_at`, and `status`. `used` (inputs) and `generated` (outputs) are present on records that represent computation (`academy_action`, node/graph records, LLM and tool calls); `parent_task_id` links child records to their enclosing parent. + +### Cross-plugin composition + +All four plugins can run under a **shared `campaign_id`** using the `from_academy_plugin()` factory. AutoGen and CrewAI share the Academy plugin's in-memory buffer directly; LangGraph creates its own interceptor but inherits the same `campaign_id`: + +```python +from flowcept.agents.academy.academy_plugin import FlowceptAcademyPlugin +from flowcept.agents.langgraph.langgraph_plugin import FlowceptLangGraphPlugin +from flowcept.agents.autogen.autogen_plugin import FlowceptAutoGenPlugin +from flowcept.agents.crewai.crewai_plugin import FlowceptCrewAIPlugin + +ap = FlowceptAcademyPlugin(config={"workflow_name": "my-run"}) +ap.start() +lg = FlowceptLangGraphPlugin.from_academy_plugin(ap) # own interceptor, shared campaign_id +ag = FlowceptAutoGenPlugin.from_academy_plugin(ap) # shared buffer +cr = FlowceptCrewAIPlugin.from_academy_plugin(ap) # shared buffer + +# ... run workloads ... + +lg.stop() # flush LangGraph's interceptor +ap.stop() # flush Academy / AutoGen / CrewAI shared buffer +``` + +Every record across all four plugins carries the same `campaign_id`, so a single query retrieves the full cross-framework provenance trace. + +### Cross-framework provenance linking + +When an Academy `@action` launches a LangGraph graph or an AutoGen team, the plugins can record an explicit parent–child edge across framework boundaries. + +**Academy (source side)** — two `contextvars.ContextVar` values are set automatically during execution: + +- `_current_academy_agent_id` — the Academy agent identifier, set once per agent at startup +- `_current_action_task_id` — the Flowcept `task_id` of the currently-executing `@action` or `@loop` + +Read the latter inside an action to get the enclosing task's identifier: + +```python +from flowcept.agents.academy.academy_plugin import _current_action_task_id +action_task_id = _current_action_task_id.get() +``` + +**LangGraph (target side)** — pass the action's `task_id` as `_source_agent_id` in the initial graph state: + +```python +result = await graph.ainvoke( + {"_source_agent_id": action_task_id, ...}, + config={"callbacks": [lg.callback_handler]}, +) +``` + +The LangGraph plugin stores it as `source_agent_id` in `custom_metadata` of both `langgraph_graph` and `langgraph_node` records. + +**AutoGen (target side)** — pass it as `source_agent_id` to `run_team()`: + +```python +result = await ag.run_team(team, task, source_agent_id=action_task_id) +``` + +The AutoGen plugin stores it in `custom_metadata` of the `autogen_run` record. + +**AI coding harness (target side)** — the harness capture (`flowcept.agents.harness`) links the other way too: give it a framework-emitted task or agent id and every turn, tool, and LLM-call task it records carries it as a top-level `source_agent_id`. Set it with the `flowcept_source_agent_id` hook-payload key, the `FLOWCEPT_HARNESS_SOURCE_AGENT_ID` environment variable (the payload key wins), or `SessionTracer(..., source_agent_id=...)`: + +```python +from flowcept.agents.harness import SessionTracer + +with SessionTracer("my_agent", source_agent_id=action_task_id) as tracer: + ... +``` + +In all cases, `campaign_id` and `workflow_id` alone are sufficient for coarse-grained cross-framework queries without explicit identifier threading. + +### Examples + +Runnable examples for each framework are in [`examples/agents/`](examples/agents/): + +- [`examples/agents/academy/academy_example.py`](examples/agents/academy/academy_example.py) +- [`examples/agents/langgraph/langgraph_example.py`](examples/agents/langgraph/langgraph_example.py) +- [`examples/agents/crewai/crewai_example.py`](examples/agents/crewai/crewai_example.py) +- [`examples/agents/autogen/autogen_example.py`](examples/agents/autogen/autogen_example.py) +- [`examples/agents/combined_agentic_systems/combined_example.py`](examples/agents/combined_agentic_systems/combined_example.py) — all four frameworks running concurrently +- [`examples/agents/langchain/langchain_example.py`](examples/agents/langchain/langchain_example.py) — LangChain callback handler capture with a fake chat model (runs offline) +- [`examples/agents/openai_agents/openai_agents_example.py`](examples/agents/openai_agents/openai_agents_example.py) — OpenAI Agents SDK tracing processor; runs a real agent with `OPENAI_API_KEY`, or synthetic SDK spans offline +- [`examples/agents/otel/otel_example.py`](examples/agents/otel/otel_example.py) — OpenTelemetry GenAI span exporter fed by synthetic spans (runs offline) +- [`examples/agents/cli_harness/cli_harness_example.py`](examples/agents/cli_harness/cli_harness_example.py) — profile-driven CLI-harness adapter, replaying Codex-style hook events (runs offline) +- [`examples/agents/claude_code/claude_code_example.py`](examples/agents/claude_code/claude_code_example.py) — Claude Code install walkthrough plus a simulated hook-event replay (runs offline) +- [`examples/agents/claude_agent_sdk/claude_agent_sdk_example.py`](examples/agents/claude_agent_sdk/claude_agent_sdk_example.py) — `trace_query` drop-in for `claude_agent_sdk.query` (requires `ANTHROPIC_API_KEY`) +- [`examples/agents/prov_analysis/prov_analysis_example.py`](examples/agents/prov_analysis/prov_analysis_example.py) — agentic provenance analysis: replays two synthetic sessions and runs every `prov_analysis` function over them (runs offline) + +## AI Coding Harness Provenance Plugins + +Flowcept also captures what an **AI coding harness** actually did — prompts, turns, +tool calls, subagents — as PROV-AGENT provenance. An agentic coding session is a +workflow: a prompt causes a turn, a turn causes tool calls, a tool call edits a +file. These plugins write that structure into Flowcept's own record format, so +"which prompt produced this bad edit?" becomes a query instead of a scroll +through a transcript. + +| Source | Plugin | +| --- | --- | +| Claude Code | [`plugins/flowcept`](plugins/flowcept) (Claude Code plugin), or hooks in `settings.json` | +| Codex CLI, Gemini CLI, Cursor, OpenCode | [`flowcept.agents.cli_harness`](src/flowcept/agents/cli_harness/) — one JSON profile per harness | +| Anything emitting OpenTelemetry GenAI spans | [`flowcept.agents.otel`](src/flowcept/agents/otel/) span exporter (`pip install flowcept[harness_otel]`) | +| Claude Agent SDK | [`flowcept.agents.claude_agent_sdk`](src/flowcept/agents/claude_agent_sdk/) `trace_query` (`pip install flowcept[harness_claude_sdk]`) | +| OpenAI Agents SDK | [`flowcept.agents.openai_agents`](src/flowcept/agents/openai_agents/) tracing processor | +| LangChain / LangGraph | [`flowcept.agents.langchain`](src/flowcept/agents/langchain/) callback handler | +| Your own agent | `flowcept.agents.harness.SessionTracer`, or the `flowcept-harness-mcp` server's `record_event` tool | + +The capture path is deliberately **stdlib-only**: a harness hook is a fresh process +on the interactive critical path, and importing heavy dependencies costs far more +than the capture itself. Records are appended as JSONL, one file per session, under +`~/.flowcept/harness/buffers/`, in Flowcept's native format. + +### Quick start: Claude Code + +``` +/plugin marketplace add +/plugin install flowcept +``` + +Then work normally, and when you want to see what was recorded: + +```bash +flowcept-harness sessions # every captured session, newest first +flowcept-harness show # the most recent one, turn by turn +flowcept-harness report # a Flowcept workflow card +``` + +### Quick start: another CLI harness + +Point the harness's hook at the profile-driven adapter: + +```bash +flowcept-harness hook --harness codex --profile codex +``` + +Profiles live in [`src/flowcept/agents/cli_harness/profiles/`](src/flowcept/agents/cli_harness/profiles/) +and are plain JSON — adding a harness means adding a file, not writing code. + +### Quick start: in-process capture + +```python +# OpenTelemetry GenAI spans +from flowcept.agents.otel.otel_plugin import FlowceptSpanExporter + +# OpenAI Agents SDK — register once, nothing else changes +from flowcept.agents.openai_agents.openai_agents_plugin import install +install() + +# LangChain / LangGraph +from flowcept.agents.langchain.langchain_plugin import FlowceptCallbackHandler + +# Claude Agent SDK — a drop-in for claude_agent_sdk.query +from flowcept.agents.claude_agent_sdk.claude_agent_sdk_plugin import trace_query + +# Your own agent +from flowcept.agents.harness import SessionTracer +``` + +See [`examples/agents/harness/harness_example.py`](examples/agents/harness/harness_example.py) +and `flowcept-harness --help` for the full CLI (status, flush to a live Flowcept +backend, repair of crashed sessions). Configuration is via `FLOWCEPT_HARNESS_*` +environment variables, including redaction of credential-shaped values, prompt +digests instead of full prompts, and offline-first buffering. Full documentation, +including the configuration table, privacy posture, and the record model, is in +[`src/flowcept/agents/harness/README.md`](src/flowcept/agents/harness/README.md). + +### Provenance analysis + +Captured sessions can be analyzed, not just replayed. The analysis logic lives +in [`flowcept.agents.prov_analysis`](src/flowcept/agents/prov_analysis/) — +pure functions over provenance records — and is exposed on every surface: + +- **Harness MCP tools** — the `flowcept-harness-mcp` server adds + `analyze_session` (summary + agent behavior), `analyze_errors` (failure + clustering with excerpts), `find_slowest` (latency ranking), and + `cross_links` (edges across framework boundaries). +- **CLI** — `flowcept-harness analyze [--errors | --slowest N | --links]` + runs the same analyses over a session's buffer. +- **Claude Code plugin** — [`plugins/flowcept`](plugins/flowcept) wires the MCP + server in via its `.mcp.json` (a stdio server launched by + `scripts/mcp-server.sh`) and ships two skills: `prov-analysis` (turn captured + provenance into answers) and `write-flowcept-plugin` (author a capture plugin + for a new harness). Set `FLOWCEPT_HARNESS_AUTOREPORT=1` to have the plugin + write a workflow card per session on SessionEnd (opt-in, off by default). +- **Flowcept agent MCP / web chat** — `df_*`/`db_*` analysis tools and + `compare_executions` over the agent's in-memory context or the DB. + +See [`examples/agents/prov_analysis/prov_analysis_example.py`](examples/agents/prov_analysis/prov_analysis_example.py) +and [`src/flowcept/agents/prov_analysis/README.md`](src/flowcept/agents/prov_analysis/README.md). + ## Storage And Querying Flowcept can run fully offline or as an online distributed system. diff --git a/deployment/compose-diaspora.yml b/deployment/compose-diaspora.yml new file mode 100644 index 00000000..5cdf0972 --- /dev/null +++ b/deployment/compose-diaspora.yml @@ -0,0 +1,44 @@ +version: '3.8' +name: flowcept +services: + # flowcept_redis: + # container_name: flowcept_redis + # image: redis + # ports: + # - 6379:6379 + + # flowcept_mongo: + # container_name: flowcept_mongo + # image: mongo:latest + # # volumes: + # # - /Users/rsr/Downloads/mongo_data/db:/data/db + # ports: + # - 27017:27017 + + # mofka: + # image: ghcr.io/mochi-hpc/mochi-spack-buildcache:mofka-0.4.0-cmuy7qp44yxutafxseqiqbn3iejima4k.spack + # ports: + # - '9999:9999' + # volumes: + # - ./resources/mofka_config.json:/config/mofka_config.json + # - ./resources/mofka.json:/config/mofka.json + + diaspora: + container_name: flowcept_diaspora + image: + ports: + - '9999:9999' + volumes: + - ./resources/diaspora_config.json:/config/diaspora_config.json + - ./resources/diaspora.json:/config/diaspora.json + entrypoint: [ '/bin/sh', '-c' ] + command: | + " + diaspora-ctl topic create --name flowcept \ + --driver files \ + --driver.root_path /tmp/diaspora-data/interception \ + --topic.num_partitions 1 + sleep 0.3 + echo "Created topic." + while true; do sleep 3600; done + " diff --git a/docs/agent.rst b/docs/agent.rst index 52be516f..a4e4fe6a 100644 --- a/docs/agent.rst +++ b/docs/agent.rst @@ -88,6 +88,40 @@ The agent resolves the matching task(s) via a Mongo-style filter, then the Dataf tab dims all unrelated nodes and edges, tracing only the ancestor/descendant chain. Click any node or empty space to reset the highlight manually. +Provenance analysis tools +------------------------- + +Both surfaces expose ready-made analysis tools built on +``flowcept.agents.prov_analysis`` (pure functions over provenance records, +shared with the coding-harness surfaces — see :doc:`harness_plugins`). Each +analysis comes in two variants: ``df_*`` tools run over the records loaded in +the agent's in-memory context (the same context the DataFrame queries use), +and ``db_*`` variants pull records from the database via ``DBAPI``, optionally +scoped by ``workflow_id``. + +.. list-table:: + :header-rows: 1 + + * - Tool + - What it returns + * - ``df_summarize_execution`` / ``db_summarize_execution`` + - Counts by activity/subtype, statuses, duration bounds, token usage. + * - ``df_analyze_errors`` / ``db_analyze_errors`` + - Per-activity error rates with stderr/message excerpts. + * - ``df_agent_behavior`` / ``db_agent_behavior`` + - Per-agent turns, tool calls, LLM calls, token usage, durations. + * - ``df_find_slowest`` / ``db_find_slowest`` + - Slowest tasks, longest elapsed first (``limit`` defaults to 10). + * - ``df_cross_framework_links`` / ``db_cross_framework_links`` + - Cross-framework provenance edges (``source_agent_id`` pointers). + * - ``compare_executions`` + - Per-activity count/duration/error-rate deltas between two workflows + (``workflow_id_a``, ``workflow_id_b``); prefers in-memory records, + falls back to the DB. + +The web chat exposes the same tools, routed to the ``df_`` or ``db_`` variant +by the chat's tool context. + Explicit MCP tool example ------------------------- diff --git a/docs/agent_plugins.rst b/docs/agent_plugins.rst new file mode 100644 index 00000000..ee375eab --- /dev/null +++ b/docs/agent_plugins.rst @@ -0,0 +1,318 @@ +Agentic Framework Provenance Plugins +==================================== + +Flowcept ships zero-code-change provenance plugins for popular agentic +frameworks: **Academy**, **AutoGen**, **CrewAI**, **LangChain**, **LangGraph**, +and the **OpenAI Agents SDK**. Each plugin automatically captures: + +- **Intra-agent provenance** — individual action/task executions with inputs, + outputs, timing, and status. +- **Inter-agent provenance** — parent/child relationships between agents and + the tasks they spawn. +- **LLM call provenance** — every OpenAI or Anthropic API call linked back to + the agent action that triggered it (model, prompt, tokens, latency). + +The Academy, AutoGen, CrewAI, and LangGraph plugins emit Flowcept +``WorkflowObject`` / ``TaskObject`` records directly and can be auto-started +from ``settings.yaml``. The LangChain callback handler and the OpenAI Agents +SDK tracing processor are in-process capture plugins that share the harness +record model described in :doc:`harness_plugins`. + +Enabling plugins via ``settings.yaml`` +-------------------------------------- + +Add a ``plugins:`` block to your ``~/.flowcept/settings.yaml``. Only the +frameworks you want to track need to be listed. Supported ``kind`` values are +``academy``, ``langgraph``, ``crewai``, and ``autogen``: + +.. code-block:: yaml + + plugins: + academy: + enabled: true + kind: academy + workflow_name: "my-academy-workflow" + performance_tracking: true + langgraph: + enabled: true + kind: langgraph + workflow_name: "my-langgraph-workflow" + performance_tracking: true + crewai: + enabled: true + kind: crewai + workflow_name: "my-crewai-workflow" + performance_tracking: true + autogen: + enabled: true + kind: autogen + workflow_name: "my-autogen-workflow" + performance_tracking: true + +Then wrap your code with ``Flowcept()`` — all enabled plugins start and stop +automatically, and every auto-started plugin inherits Flowcept's +``campaign_id``: + +.. code-block:: python + + from flowcept import Flowcept + + with Flowcept(): + # your Academy / LangGraph / CrewAI / AutoGen code here + ... + +Running plugin instances are available as ``flowcept_instance.plugins``, a +dict keyed by the plugin's config name. + +Helper utilities +---------------- + +The Academy, AutoGen, CrewAI, and LangGraph plugin modules each export drop-in +wrappers that record LLM calls regardless of which framework is active: + +- ``openai_chat(prompt, model, ...)`` — call OpenAI and automatically record + the call with provenance linkage. +- ``anthropic_chat(prompt, model, ...)`` — the same for Anthropic Claude + models. +- ``FlowceptAnthropicClient(client)`` — wrap an existing + ``anthropic.Anthropic`` client to intercept all ``messages.create`` / + ``stream`` calls. +- ``run_team(team, task, ...)`` *(AutoGen only)* — run an AutoGen team and + capture full provenance. + +For example: ``from flowcept.agents.academy.academy_plugin import openai_chat``. + +Per-plugin reference +-------------------- + +Academy +~~~~~~~ + +``flowcept.agents.academy.academy_plugin.FlowceptAcademyPlugin`` wraps Academy +agents automatically and records ``@action`` and ``@loop`` executions. + +Record hierarchy: campaign ``WorkflowObject`` → agent ``WorkflowObject`` → +``academy_action`` / ``academy_loop`` / ``academy_lifecycle`` task records +(siblings), with ``llm_call`` records nested under the action or loop that +issued them. + +Two ``contextvars.ContextVar`` values are set automatically during execution +and can be read from user code: + +- ``_current_academy_agent_id`` — the Academy agent identifier, set once per + agent at startup. +- ``_current_action_task_id`` — the Flowcept ``task_id`` of the + currently-executing ``@action`` or ``@loop``. + +AutoGen +~~~~~~~ + +``flowcept.agents.autogen.autogen_plugin.FlowceptAutoGenPlugin`` wraps AutoGen +teams automatically. The module also exports ``run_team(team, task, ...)``, +which runs a team and captures full provenance, and ``FlowceptModelClient``, +a model-client wrapper. + +Record hierarchy: ``WorkflowObject`` → ``autogen_run`` → ``autogen_message`` +→ ``llm_call``. + +CrewAI +~~~~~~ + +``flowcept.agents.crewai.crewai_plugin.FlowceptCrewAIPlugin`` records crew, +task, and agent executions through CrewAI's event listener and hook +interfaces (no automatic agent wrapping is needed or performed). + +Record hierarchy: ``WorkflowObject`` → ``crewai_crew`` (no children via +``parent_task_id``); ``crewai_task`` → ``crewai_agent`` → ``llm_call`` / +``tool_call``. + +LangChain +~~~~~~~~~ + +``flowcept.agents.langchain.langchain_plugin.FlowceptCallbackHandler`` is a +LangChain callback handler that records each root chain/graph run as a turn, +plus the model and tool calls inside it. Pass it wherever LangChain accepts +callbacks: + +.. code-block:: python + + from flowcept.agents.langchain.langchain_plugin import FlowceptCallbackHandler + + graph.invoke(state, config={"callbacks": [FlowceptCallbackHandler(session_id="thread-42")]}) + +Records go to the harness session buffer (see :doc:`harness_plugins` for the +record model and where records are written). + +LangGraph +~~~~~~~~~ + +``flowcept.agents.langgraph.langgraph_plugin.FlowceptLangGraphPlugin`` records +graph and node executions through a LangGraph/LangChain callback exposed as +``plugin.callback_handler``: + +.. code-block:: python + + result = graph.invoke(state, config={"callbacks": [plugin.callback_handler]}) + +Record hierarchy: ``WorkflowObject`` → ``langgraph_graph`` → +``langgraph_node`` → ``llm_call`` / ``tool_call``. + +OpenAI Agents SDK +~~~~~~~~~~~~~~~~~ + +``flowcept.agents.openai_agents.openai_agents_plugin`` provides +``FlowceptTraceProcessor``, a tracing processor for the OpenAI Agents SDK. +Register it once and nothing else changes: + +.. code-block:: python + + from flowcept.agents.openai_agents.openai_agents_plugin import install + + install() + +Like the LangChain handler, it records into the harness session buffer +(see :doc:`harness_plugins`). + +What gets captured +------------------ + +For the four framework plugins (Academy, LangGraph, CrewAI, AutoGen): + +.. list-table:: + :header-rows: 1 + + * - Captured field + - Academy + - LangGraph + - CrewAI + - AutoGen + * - Agent action / node executions + - ✓ + - ✓ + - ✓ + - ✓ + * - Inputs and outputs per action + - ✓ + - ✓ + - ✓ + - ✓ + * - Timing (start / end / latency) + - ✓ + - ✓ + - ✓ + - ✓ + * - Parent–child task linkage + - ✓ + - ✓ + - ✓ + - ✓ + * - LLM calls (OpenAI and Anthropic) + - ✓ + - ✓ + - ✓ + - ✓ + * - Token usage + - ✓ + - ✓ + - ✓ + - ✓ + * - Agent ID on LLM calls + - ✓ + - ✓ + - ✓ + - ✓ + * - Automatic agent wrapping + - ✓ + - ✓ + - — + - ✓ + +Each plugin emits typed ``TaskObject`` records tagged with a ``subtype`` +field, forming a nested provenance hierarchy (see the per-plugin sections +above for each hierarchy). All records carry ``campaign_id``, +``workflow_id``, ``task_id``, ``started_at``, ``ended_at``, and ``status``. +``used`` (inputs) and ``generated`` (outputs) are present on records that +represent computation (``academy_action``, node/graph records, LLM and tool +calls); ``parent_task_id`` links child records to their enclosing parent. + +Cross-plugin composition +------------------------ + +The Academy, LangGraph, CrewAI, and AutoGen plugins can run under a **shared +campaign** using the ``from_academy_plugin()`` factory. AutoGen and CrewAI +share the Academy plugin's in-memory buffer directly; LangGraph creates its +own interceptor but inherits the same ``campaign_id``: + +.. code-block:: python + + from flowcept.agents.academy.academy_plugin import FlowceptAcademyPlugin + from flowcept.agents.langgraph.langgraph_plugin import FlowceptLangGraphPlugin + from flowcept.agents.autogen.autogen_plugin import FlowceptAutoGenPlugin + from flowcept.agents.crewai.crewai_plugin import FlowceptCrewAIPlugin + + ap = FlowceptAcademyPlugin(config={"workflow_name": "my-run"}) + ap.start() + lg = FlowceptLangGraphPlugin.from_academy_plugin(ap) # own interceptor, shared campaign_id + ag = FlowceptAutoGenPlugin.from_academy_plugin(ap) # shared buffer + cr = FlowceptCrewAIPlugin.from_academy_plugin(ap) # shared buffer + + # ... run workloads ... + + lg.stop() # flush LangGraph's interceptor + ap.stop() # flush Academy / AutoGen / CrewAI shared buffer + +Every record across all four plugins carries the same ``campaign_id``, so a +single query retrieves the full cross-framework provenance trace. + +Cross-framework provenance linking +---------------------------------- + +When an Academy ``@action`` launches a LangGraph graph or an AutoGen team, the +plugins can record an explicit parent–child edge across framework boundaries. + +**Academy (source side)** — read the enclosing task's identifier inside an +action: + +.. code-block:: python + + from flowcept.agents.academy.academy_plugin import _current_action_task_id + + action_task_id = _current_action_task_id.get() + +**LangGraph (target side)** — pass the action's ``task_id`` as +``_source_agent_id`` in the initial graph state: + +.. code-block:: python + + result = await graph.ainvoke( + {"_source_agent_id": action_task_id, ...}, + config={"callbacks": [lg.callback_handler]}, + ) + +The LangGraph plugin stores it as ``source_agent_id`` in ``custom_metadata`` +of both ``langgraph_graph`` and ``langgraph_node`` records. + +**AutoGen (target side)** — pass it as ``source_agent_id`` to ``run_team()``: + +.. code-block:: python + + result = await ag.run_team(team, task, source_agent_id=action_task_id) + +The AutoGen plugin stores it in ``custom_metadata`` of the ``autogen_run`` +record. + +In all cases, ``campaign_id`` and ``workflow_id`` alone are sufficient for +coarse-grained cross-framework queries without explicit identifier threading. + +Examples +-------- + +Runnable examples for each framework are in +`examples/agents/ `_: + +- ``examples/agents/academy/academy_example.py`` +- ``examples/agents/langgraph/langgraph_example.py`` +- ``examples/agents/crewai/crewai_example.py`` +- ``examples/agents/autogen/autogen_example.py`` +- ``examples/agents/combined_agentic_systems/combined_example.py`` — all four + frameworks running concurrently under one campaign diff --git a/docs/harness_plugins.rst b/docs/harness_plugins.rst new file mode 100644 index 00000000..a904ad11 --- /dev/null +++ b/docs/harness_plugins.rst @@ -0,0 +1,322 @@ +AI Coding Harness Provenance Plugins +==================================== + +Flowcept can capture what an **AI coding harness** actually did — prompts, +turns, tool calls, subagents — as PROV-AGENT provenance. An agentic coding +session is a workflow: a prompt causes a turn, a turn causes tool calls, a +tool call edits a file. These plugins write that structure into Flowcept's own +record format, so "which prompt produced this bad edit?" becomes a query +instead of a scroll through a transcript. + +The shared capture core is ``flowcept.agents.harness``; the per-source plugin +modules live beside it under ``flowcept/agents/``: + +.. list-table:: + :header-rows: 1 + + * - Source + - Plugin + * - Claude Code + - ``plugins/flowcept`` (Claude Code plugin), or hooks in ``settings.json`` + (adapter: ``flowcept.agents.claude_code``) + * - Codex CLI, Gemini CLI, Cursor, OpenCode + - ``flowcept.agents.cli_harness`` — one JSON profile per harness + * - Anything emitting OpenTelemetry GenAI spans + - ``flowcept.agents.otel`` span exporter + (``pip install "flowcept[harness_otel]"``) + * - Claude Agent SDK + - ``flowcept.agents.claude_agent_sdk`` ``trace_query`` + (``pip install "flowcept[harness_claude_sdk]"``) + * - OpenAI Agents SDK + - ``flowcept.agents.openai_agents`` tracing processor + * - LangChain / LangGraph + - ``flowcept.agents.langchain`` callback handler + * - Your own agent + - ``flowcept.agents.harness.SessionTracer``, or the + ``flowcept-harness-mcp`` server's ``record_event`` tool + +The capture path is deliberately **stdlib-only**: a harness hook is a fresh +process on the interactive critical path, and importing heavy dependencies +costs far more than the capture itself. Flowcept's heavier dependencies are +only loaded to *read* what was captured. + +Quick start: Claude Code +------------------------ + +Install the Claude Code plugin from this repository: + +.. code-block:: text + + /plugin marketplace add + /plugin install flowcept + +Then work normally. When you want to see what was recorded: + +.. code-block:: bash + + flowcept-harness sessions # every captured session, newest first + flowcept-harness show # the most recent one, turn by turn + flowcept-harness report # a Flowcept workflow card + +``flowcept-harness install`` prints the equivalent ``settings.json`` if you +would rather wire the hooks yourself than use the plugin. + +Quick start: another CLI harness +-------------------------------- + +Point the harness's hook at the generic profile-driven adapter: + +.. code-block:: bash + + flowcept-harness hook --harness codex --profile codex + +Profiles live in ``src/flowcept/agents/cli_harness/profiles/`` (``codex``, +``gemini``, ``cursor``, ``opencode``) and are plain JSON: a map from the +harness's event names to normalized ones, and a map from its payload fields +to Flowcept's. Adding a harness means adding a file, not writing code. +``--profile`` also accepts a filesystem path, so a profile can live outside +the package while you iterate on it. + +Quick start: Claude Agent SDK +----------------------------- + +``trace_query`` is a drop-in for ``claude_agent_sdk.query``: + +.. code-block:: python + + from flowcept.agents.claude_agent_sdk.claude_agent_sdk_plugin import trace_query + + async for message in trace_query(prompt="fix the failing test"): + ... + +Quick start: OpenTelemetry GenAI spans +-------------------------------------- + +In-process, for anything already instrumented with OpenTelemetry: + +.. code-block:: python + + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from flowcept.agents.otel.otel_plugin import FlowceptSpanExporter + + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(FlowceptSpanExporter())) + +Or after the fact, from spans a collector already wrote: + +.. code-block:: python + + from flowcept.agents.otel.otel_plugin import ingest_file + + ingest_file("spans.jsonl") + +Spans are read through the OTel GenAI semantic conventions: +``gen_ai.operation.name`` separates a model call from a tool call, and +``gen_ai.conversation.id`` groups spans into a session. Spans that are not +GenAI spans are ignored — an HTTP client span is not provenance. + +Quick start: in-process capture +------------------------------- + +.. code-block:: python + + # OpenAI Agents SDK — register once, nothing else changes + from flowcept.agents.openai_agents.openai_agents_plugin import install + install() + + # LangChain / LangGraph — pass the handler as a callback + from flowcept.agents.langchain.langchain_plugin import FlowceptCallbackHandler + graph.invoke(state, config={"callbacks": [FlowceptCallbackHandler(session_id="thread-42")]}) + +Each wrapper is duck-typed against its SDK — nothing imports the SDK it +wraps, so installing one does not drag in the others. + +For an agent that is none of the above, drive the session yourself: + +.. code-block:: python + + from flowcept.agents.harness import SessionTracer + + with SessionTracer("my_agent", model="claude-opus-5") as tracer: + tracer.prompt("summarize the repo") + with tracer.tool("read_file", {"path": "README.md"}) as call: + call.result(read("README.md")) + tracer.turn_end("done", usage={"input_tokens": 900}) + +What gets recorded +------------------ + +.. code-block:: text + + session workflow (subtype: agent_session) + turn task (subtype: ai_model_invocation, granularity=turn) + model call task (subtype: ai_model_invocation, granularity=call) + tool call task (subtype: agent_tool) + subagent workflow (subtype: subagent_session) + task + compaction, notice task (subtype: harness_event) + +The edges are the point. A tool call's ``parent_task_id`` is its turn; a +subagent's tools live in the subagent's own workflow rather than interleaved +with the parent's; a turn lists the tools it caused. That is the +``wasInformedBy`` chain PROV-AGENT is built around, and it is what lets you +walk from a bad edit back to the prompt that caused it. + +Model invocations are recorded at whatever granularity the source can see. A +hook cannot observe individual API calls, so hook-based capture records one +invocation per turn; SDK and OTel capture record both. ``granularity`` in +``custom_metadata`` says which you are looking at. + +The buffer / flush model +------------------------ + +Records are appended as JSONL, one file per session, under +``~/.flowcept/harness/buffers/``. The format is Flowcept's native record +format, so the buffer is directly consumable: + +.. code-block:: bash + + flowcept --generate-report --input-path ~/.flowcept/harness/buffers/.jsonl + +To push into a live Flowcept backend instead of (or as well as) the file: + +.. code-block:: bash + + export FLOWCEPT_HARNESS_ONLINE=1 # publish as you go + flowcept-harness flush --all # or publish buffers after the fact + +Offline is the default because a hook must never block on a message queue +that may not be running. + +The ``flowcept-harness`` CLI +---------------------------- + +.. code-block:: text + + flowcept-harness sessions list captured sessions, newest first + flowcept-harness show show one session's activity, turn by turn + flowcept-harness status configuration and capture health + flowcept-harness report generate a Flowcept report from a buffer + flowcept-harness analyze analyze one session (failures, latency, links) + flowcept-harness flush publish buffered records to a Flowcept backend + flowcept-harness repair close sessions a crashed harness left open + flowcept-harness install print the settings that enable capture + flowcept-harness hook record a payload from stdin (what hooks call) + +Useful options: + +- ``sessions -n/--limit N`` and ``-v/--verbose`` +- ``show [buffer ...]`` — defaults to the most recent session +- ``status --check-backend`` — also probe the Flowcept backend +- ``flush --input | --all``, plus ``--remove`` (delete buffers + after a successful flush) and ``--dry-run`` +- ``report --input --type workflow_card --format markdown + -o/--output `` +- ``--home `` (global) — override the capture home directory + +Provenance analysis +------------------- + +Captured sessions can be analyzed, not just replayed. The analysis logic is +``flowcept.agents.prov_analysis`` — pure functions over provenance records +(``src/flowcept/agents/prov_analysis/README.md``) — exposed on three +harness-side surfaces: + +**MCP analysis tools.** The ``flowcept-harness-mcp`` server exposes four +analysis tools alongside its capture and inspection tools; each takes an +optional ``session`` id prefix and defaults to the most recent session: + +- ``analyze_session`` — execution summary (counts, statuses, durations, token + usage) plus per-agent behavior (turns, tool calls, subagent fan-out). +- ``analyze_errors`` — failure clustering: which activities fail, how often, + with stderr/message excerpts. +- ``find_slowest`` — the slowest tasks, longest elapsed first, with activity, + status, and parent-chain depth (``limit`` defaults to 10). +- ``cross_links`` — cross-framework provenance edges built from + ``source_agent_id`` pointers (e.g. a LangGraph run launched from a coding + session), plus the count of unlinked tasks. The harness side of the edge is + written by giving capture a framework-emitted task/agent id: the + ``flowcept_source_agent_id`` hook-payload key, the + ``FLOWCEPT_HARNESS_SOURCE_AGENT_ID`` environment variable (payload key + wins), or ``SessionTracer(..., source_agent_id=...)``. + +**CLI.** The same analyses from the terminal: + +.. code-block:: bash + + flowcept-harness analyze # summary + agent behavior + flowcept-harness analyze --errors # failure clustering + flowcept-harness analyze --slowest 5 # latency ranking + flowcept-harness analyze --links # cross-framework links + flowcept-harness analyze --compare # per-activity deltas of two sessions + +```` is a workflow id or prefix and defaults to the most recent +session. ``--compare`` takes two ids/prefixes instead of the positional +session and prints per-activity count, average-duration, and error-rate +deltas. + +**Claude Code plugin.** ``plugins/flowcept`` wires the MCP server in through +its ``.mcp.json``, which declares a ``flowcept-provenance`` stdio server +launched by ``scripts/mcp-server.sh`` — installing the plugin gives Claude +Code the analysis tools with no extra setup. The plugin also ships two +analysis-oriented skills: ``prov-analysis`` (turn captured provenance into +answers: failures, latency, cost, run comparison, cross-framework linkage) +and ``write-flowcept-plugin`` (author a capture plugin for a new harness). +Setting ``FLOWCEPT_HARNESS_AUTOREPORT=1`` additionally enables an opt-in +SessionEnd hook that writes a workflow card per session to +``$FLOWCEPT_HARNESS_HOME/reports/``; it is off by default and exits silently +when unset. + +The Flowcept agent MCP server and web chat expose the same analyses over +in-memory and DB records — see :doc:`agent`. A runnable, fully offline +example is ``examples/agents/prov_analysis/prov_analysis_example.py``. + +Configuration +------------- + +Every knob is an environment variable, so it can be set from a harness +settings file, a plugin's ``userConfig``, or a shell profile. The most +important ones: + +.. list-table:: + :header-rows: 1 + + * - Variable + - Default + - Meaning + * - ``FLOWCEPT_HARNESS_ENABLED`` + - ``1`` + - Master switch. + * - ``FLOWCEPT_HARNESS_HOME`` + - ``~/.flowcept/harness`` + - State and buffers. + * - ``FLOWCEPT_HARNESS_CONTENT`` + - ``summary`` + - File bodies: ``full``, ``summary``, ``none``. + * - ``FLOWCEPT_HARNESS_REDACT`` + - ``1`` + - Redact credential-shaped keys and literals. + * - ``FLOWCEPT_HARNESS_CAPTURE_PROMPTS`` + - ``1`` + - Off stores prompt digests only. + * - ``FLOWCEPT_HARNESS_ONLINE`` + - ``0`` + - Publish to the Flowcept MQ as you go. + * - ``FLOWCEPT_HARNESS_TIMEOUT_MS`` + - ``2000`` + - Hard ceiling on hook wall time. + +The full configuration table (campaign scoping, string truncation, tool-result +capture, debug logging), the privacy posture, and the MCP server +(``flowcept-harness-mcp``) are documented in the package README: +`src/flowcept/agents/harness/README.md +`_. + +See also +-------- + +- ``examples/agents/harness/harness_example.py`` — a ``SessionTracer`` + example. +- :doc:`agent_plugins` — provenance plugins for agentic frameworks + (Academy, AutoGen, CrewAI, LangChain, LangGraph, OpenAI Agents SDK). +- :doc:`schemas` — the PROV-AGENT data model in Flowcept. diff --git a/docs/index.rst b/docs/index.rst index b0ec9be0..4a785fff 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -65,6 +65,8 @@ Flowcept setup web_ui agent + agent_plugins + harness_plugins prov_capture telemetry_capture prov_storage diff --git a/docs/openapi/flowcept-openapi.json b/docs/openapi/flowcept-openapi.json index c665e52c..208f4df4 100644 --- a/docs/openapi/flowcept-openapi.json +++ b/docs/openapi/flowcept-openapi.json @@ -3202,6 +3202,29 @@ } } } + }, + "/": { + "get": { + "tags": [ + "health" + ], + "summary": "Root", + "operationId": "root__get", + "responses": { + "200": { + "description": "Successful Response", + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "type": "object", + "title": "Response Root Get" + } + } + } + } + } + } } }, "components": { @@ -3890,6 +3913,13 @@ "type": { "type": "string", "title": "Error Type" + }, + "input": { + "title": "Input" + }, + "ctx": { + "type": "object", + "title": "Context" } }, "type": "object", diff --git a/docs/openapi/flowcept-openapi.yaml b/docs/openapi/flowcept-openapi.yaml index 07c0d552..4d3cf0be 100644 --- a/docs/openapi/flowcept-openapi.yaml +++ b/docs/openapi/flowcept-openapi.yaml @@ -2006,6 +2006,21 @@ paths: application/json: schema: $ref: '#/components/schemas/HTTPValidationError' + /: + get: + tags: + - health + summary: Root + operationId: root__get + responses: + '200': + description: Successful Response + content: + application/json: + schema: + additionalProperties: true + type: object + title: Response Root Get components: schemas: AggregationSpec: @@ -2457,6 +2472,11 @@ components: type: type: string title: Error Type + input: + title: Input + ctx: + type: object + title: Context type: object required: - loc diff --git a/examples/agents/academy/academy_example.py b/examples/agents/academy/academy_example.py new file mode 100644 index 00000000..d7b86741 --- /dev/null +++ b/examples/agents/academy/academy_example.py @@ -0,0 +1,172 @@ +""" +Academy example +=============== + +Two Academy agents increment a counter five times (1+2+3+4+5 = 15), then an +llm_chat() call (OpenAI or Anthropic) interprets the final value — mirroring the AutoGen / CrewAI / +LangGraph examples. + +Run +--- + OPENAI_API_KEY=sk-... (or ANTHROPIC_API_KEY=sk-ant-...) python examples/agents/academy/academy_example.py + +Plugin configuration +-------------------- +Enable the Academy plugin in your settings.yaml: + + plugins: + academy: + enabled: true + kind: academy + workflow_name: "academy-counter" + performance_tracking: true + +Flowcept will auto-start/stop the plugin — no explicit plugin.start() / +plugin.stop() calls needed. +""" +from __future__ import annotations + +import asyncio +import os +import sys +from concurrent.futures import ThreadPoolExecutor + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + +from academy.agent import Agent, action, loop +from academy.exchange import LocalExchangeFactory +from academy.manager import Manager + +from flowcept import Flowcept +# Use whichever LLM provider has a key in the environment. +if os.getenv("ANTHROPIC_API_KEY"): + from flowcept.agents.academy.academy_plugin import anthropic_chat as llm_chat + + LLM_MODEL = "claude-haiku-4-5-20251001" +else: + from flowcept.agents.academy.academy_plugin import openai_chat as llm_chat + + LLM_MODEL = "gpt-4o-mini" + + +# --------------------------------------------------------------------------- +# Agent definitions +# --------------------------------------------------------------------------- + + +class CounterAgent(Agent): + """Simple integer counter.""" + + def __init__(self) -> None: + super().__init__() + self._value: int = 0 + + @action + async def increment(self, n: int = 1) -> int: + self._value += n + return self._value + + @action + async def get_value(self) -> int: + return self._value + + +class SummaryAgent(Agent): + """Increments the CounterAgent five times and records the final value.""" + + def __init__(self) -> None: + super().__init__() + self._counter = None + self._result: int = 0 + self._done: bool = False + + @action + async def set_counter(self, counter) -> None: + self._counter = counter + + @action + async def get_result(self) -> int: + return self._result + + @action + async def is_done(self) -> bool: + return self._done + + @loop + async def summary_loop(self, shutdown: asyncio.Event) -> None: + while self._counter is None and not shutdown.is_set(): + await asyncio.sleep(0.05) + if shutdown.is_set(): + return + + print("[SummaryAgent] Starting counter cycle …", flush=True) + for step in range(1, 6): + value = await self._counter.increment(step) + print(f"[SummaryAgent] increment({step}) → counter={value}", flush=True) + + self._result = await self._counter.get_value() + print(f"[SummaryAgent] Final counter value: {self._result}", flush=True) + self._done = True + + +# --------------------------------------------------------------------------- +# Driver +# --------------------------------------------------------------------------- + + +async def _run() -> int: + exchange = LocalExchangeFactory() + executor = ThreadPoolExecutor(max_workers=4) + + async with await Manager.from_exchange_factory( + factory=exchange, + executors=executor, + ) as manager: + counter = await manager.launch(CounterAgent) + await counter.ping() + + summary = await manager.launch(SummaryAgent) + await summary.ping() + + await summary.set_counter(counter) + + for _ in range(60): + await asyncio.sleep(0.5) + if await summary.is_done(): + break + + return await summary.get_result() + + +def main() -> None: + import logging + + logging.basicConfig(level=logging.INFO) + + with Flowcept(): + result = asyncio.run(_run()) + + print("\n" + "=" * 60) + print("RESULT") + print("=" * 60) + print(f"\nCounter reached {result} (expected 15 = 1+2+3+4+5)", flush=True) + assert result == 15, f"Expected 15 but got {result}" + + print("\n[example] Calling llm_chat() to interpret the counter result …", flush=True) + interpretation = llm_chat( + prompt=( + f"A counter was incremented 5 times with values 1, 2, 3, 4, 5 " + f"and reached a final value of {result}. " + f"In exactly one sentence, explain what this arithmetic result represents." + ), + model=LLM_MODEL, + system="You are a concise data analyst.", + temperature=0.3, + context={"agent": "SummaryAgent", "call_type": "interpret_result"}, + ) + print(f"\n[example] LLM says: {interpretation}\n", flush=True) + print("[example] Assertion passed.", flush=True) + + +if __name__ == "__main__": + main() diff --git a/examples/agents/autogen/autogen_example.py b/examples/agents/autogen/autogen_example.py new file mode 100644 index 00000000..7d804267 --- /dev/null +++ b/examples/agents/autogen/autogen_example.py @@ -0,0 +1,220 @@ +""" +examples/agents/autogen/autogen_example.py +========================================== + +Counter test mirroring the Academy example: + + AssistantAgent : increments a counter 5 times (values 1-5) + CriticAgent : verifies each running total + interpret step : calls llm_chat() inside _run() to interpret the final + value (15 = 1+2+3+4+5), mirroring SummaryAgent.interpret_result() + +Run +--- + OPENAI_API_KEY=sk-... (or ANTHROPIC_API_KEY=sk-ant-...) python examples/agents/autogen/autogen_example.py + +Plugin configuration +-------------------- +Enable the AutoGen plugin in your settings.yaml: + + plugins: + autogen: + enabled: true + kind: autogen + workflow_name: "autogen-counter-test" + performance_tracking: true + +Flowcept will auto-start/stop the plugin — no explicit plugin.start() / +plugin.stop() calls needed. +""" +from __future__ import annotations + +import asyncio +import os +import sys +from typing import AsyncGenerator + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + +from flowcept import Flowcept +# Use whichever LLM provider has a key in the environment. +if os.getenv("ANTHROPIC_API_KEY"): + from flowcept.agents.autogen.autogen_plugin import anthropic_chat as llm_chat + + LLM_MODEL = "claude-haiku-4-5-20251001" +else: + from flowcept.agents.autogen.autogen_plugin import openai_chat as llm_chat + + LLM_MODEL = "gpt-4o-mini" +from flowcept.agents.autogen.autogen_plugin import run_team + +try: + from autogen_agentchat.agents import AssistantAgent + from autogen_agentchat.conditions import MaxMessageTermination + from autogen_agentchat.teams import RoundRobinGroupChat + from autogen_core.models import ChatCompletionClient, CreateResult, RequestUsage, ModelInfo +except ImportError: + print("ERROR: pip install autogen-agentchat", file=sys.stderr) + sys.exit(1) + + +# --------------------------------------------------------------------------- +# Stub model client — canned counter responses, no API call needed +# --------------------------------------------------------------------------- + +_STEPS = [1, 2, 3, 4, 5] +_RUNNING = 0 +_STUB_RESPONSES = [] +for _s in _STEPS: + _RUNNING += _s + _STUB_RESPONSES.append( + f"assistant: increment({_s}) → running total = {_RUNNING}" + ) +_STUB_RESPONSES.append( + f"critic: all 5 increments verified. Final value = {_RUNNING}. TERMINATE" +) + +_stub_idx = 0 + + +class _StubModelClient(ChatCompletionClient): + """Returns canned counter responses without a real API call.""" + + @property + def model_info(self): + return ModelInfo(vision=False, function_calling=False, json_output=False, + family="stub", structured_output=False) + + @property + def capabilities(self): + return ModelInfo(vision=False, function_calling=False, json_output=False) + + async def create(self, *_, **__) -> CreateResult: + global _stub_idx + content = _STUB_RESPONSES[_stub_idx % len(_STUB_RESPONSES)] + _stub_idx += 1 + return CreateResult( + content=content, + usage=RequestUsage(prompt_tokens=20, completion_tokens=20), + finish_reason="stop", + cached=False, + ) + + async def create_stream(self, *_, **__) -> AsyncGenerator: + yield await self.create() + + def actual_usage(self) -> RequestUsage: + return RequestUsage(prompt_tokens=0, completion_tokens=0) + + def total_usage(self) -> RequestUsage: + return RequestUsage(prompt_tokens=0, completion_tokens=0) + + def count_tokens(self, *_, **__) -> int: + return 0 + + def remaining_tokens(self, *_, **__) -> int: + return 4096 + + async def close(self) -> None: + pass + + +# --------------------------------------------------------------------------- +# Agents and team +# --------------------------------------------------------------------------- + +class AssistantCounterAgent(AssistantAgent): + """Increments a counter and reports the running total.""" + + def __init__(self): + super().__init__( + name="assistant", + model_client=_StubModelClient(), + system_message=( + "You are a counter agent. Increment the counter by the given step " + "and report the running total. When all 5 steps are done, end with TERMINATE." + ), + ) + + +class CriticAgent(AssistantAgent): + """Verifies each running total.""" + + def __init__(self): + super().__init__( + name="critic", + model_client=_StubModelClient(), + system_message=( + "You are a verifier. Check each running total and confirm it is correct. " + "When TERMINATE is reached, echo it." + ), + ) + + +# --------------------------------------------------------------------------- +# Driver — mirrors academy_example._run() +# --------------------------------------------------------------------------- + +async def _run() -> None: + assistant = AssistantCounterAgent() + critic = CriticAgent() + termination = MaxMessageTermination(max_messages=len(_STUB_RESPONSES) + 1) + team = RoundRobinGroupChat( + participants=[assistant, critic], + termination_condition=termination, + ) + + task = "Increment a counter 5 times with step values 1, 2, 3, 4, 5 and report each running total." + print(f"\n[example] Task: {task!r}\n", flush=True) + + # run_team() uses _ACTIVE_INTERCEPTOR — no explicit plugin handle needed, + # mirroring how openai_chat() / anthropic_chat() work. + result = await run_team(team, task, team_name="counter-team") + + # Parse final total — mirrors SummaryAgent.get_result() + final_value = 0 + for msg in result.messages: + content = getattr(msg, "content", "") + if isinstance(content, str): + for token in content.split(): + try: + final_value = int(token.strip(".,")) + except ValueError: + pass + + print("\n" + "=" * 60) + print("TEAM CONVERSATION") + print("=" * 60) + for msg in result.messages: + source = getattr(msg, "source", "?") + content = getattr(msg, "content", "") + print(f"\n[{source}]\n{content}") + print(f"\nStop reason: {result.stop_reason}") + + # Interpret the result — mirrors SummaryAgent.interpret_result() @action. + # llm_chat() fires record_llm_call() automatically. + print("\n[example] Calling llm_chat() to interpret the counter result …", flush=True) + interpretation = llm_chat( + prompt=( + f"A counter was incremented 5 times with values 1, 2, 3, 4, 5 " + f"and reached a final value of {final_value}. " + f"In exactly one sentence, explain what this arithmetic result represents." + ), + model=LLM_MODEL, + system="You are a concise data analyst.", + temperature=0.3, + context={"agent": "counter-team", "call_type": "interpret_result"}, + ) + print(f"\n[example] LLM says: {interpretation}\n", flush=True) + print(f"[example] Counter reached {final_value} (expected 15 = 1+2+3+4+5)", flush=True) + assert final_value == 15, f"Expected 15 but got {final_value}" + print("[example] Assertion passed.", flush=True) + + +def main() -> None: + with Flowcept(): + asyncio.run(_run()) + + +if __name__ == "__main__": + main() diff --git a/examples/agents/claude_agent_sdk/claude_agent_sdk_example.py b/examples/agents/claude_agent_sdk/claude_agent_sdk_example.py new file mode 100644 index 00000000..4c23352c --- /dev/null +++ b/examples/agents/claude_agent_sdk/claude_agent_sdk_example.py @@ -0,0 +1,109 @@ +""" +Claude Agent SDK provenance capture through trace_query. + +``trace_query`` is a drop-in for ``claude_agent_sdk.query``: same arguments, +same yielded messages, provenance as a side effect. The SDK hands you an async +stream of message objects, and everything provenance needs — tool uses, tool +results, token usage, the final answer — is already in that stream, so +swapping ``query`` for ``trace_query`` is the entire integration. (To capture +a ``ClaudeSDKClient`` conversation instead, drive ``ClaudeAgentTracer`` +directly with the messages you receive.) + +The example asks one question with a single-turn budget and no tools, prints +the streamed answer, and leaves the session as PROV-AGENT provenance in a +JSONL buffer under ``~/.flowcept/harness/buffers/``. + +Run +--- + ANTHROPIC_API_KEY=sk-ant-... python examples/agents/claude_agent_sdk/claude_agent_sdk_example.py + +Requires the Claude Agent SDK: ``pip install "flowcept[harness_claude_sdk]"`` +(or ``pip install claude-agent-sdk``). + +Then inspect the capture: + + flowcept-harness sessions + flowcept-harness show + flowcept-harness report + +Plugin configuration +-------------------- +The harness plugins are configured with environment variables, not +settings.yaml (see src/flowcept/agents/harness/README.md for the full table): + + FLOWCEPT_HARNESS_ENABLED=1 # master switch (default) + FLOWCEPT_HARNESS_ONLINE=0 # 1 publishes to a live Flowcept backend + FLOWCEPT_HARNESS_REDACT=1 # redact credential-shaped values +""" + +from __future__ import annotations + +import asyncio +import os +import sys + +from flowcept.agents.claude_agent_sdk.claude_agent_sdk_plugin import trace_query + +try: + from claude_agent_sdk import AssistantMessage, ClaudeAgentOptions, CLINotFoundError, ResultMessage +except ImportError: + print('ERROR: pip install "flowcept[harness_claude_sdk]" (or: pip install claude-agent-sdk)', file=sys.stderr) + sys.exit(1) + + +async def run() -> None: + """Ask one question through trace_query and print the streamed answer.""" + options = ClaudeAgentOptions( + max_turns=1, + allowed_tools=[], # a pure Q&A turn; tool use would be captured too + system_prompt="You are a concise data analyst.", + ) + + prompt = ( + "A counter was incremented 5 times with values 1, 2, 3, 4, 5 and reached 15. " + "In exactly one sentence, explain what this arithmetic result represents." + ) + print(f"\n[example] Prompt: {prompt}\n", flush=True) + + # trace_query yields exactly what claude_agent_sdk.query yields; the + # provenance capture is a side effect. + async for message in trace_query(prompt=prompt, options=options): + if isinstance(message, AssistantMessage): + for block in message.content: + text = getattr(block, "text", None) + if isinstance(text, str): + print(f"[example] Claude: {text}", flush=True) + elif isinstance(message, ResultMessage): + print(f"\n[example] Turns: {message.num_turns}, cost: ${message.total_cost_usd or 0:.4f}", flush=True) + + +def main(): + """Check credentials, then run the traced query.""" + if not os.getenv("ANTHROPIC_API_KEY"): + print( + "ERROR: ANTHROPIC_API_KEY is not set.\n" + "This example calls the Claude API through the Claude Agent SDK.\n" + "Set the key and re-run:\n" + " ANTHROPIC_API_KEY=sk-ant-... python examples/agents/claude_agent_sdk/claude_agent_sdk_example.py", + file=sys.stderr, + ) + sys.exit(1) + + try: + asyncio.run(run()) + except CLINotFoundError: + print( + "ERROR: the Claude Code CLI the SDK drives was not found.\n" + "Install it with: npm install -g @anthropic-ai/claude-code", + file=sys.stderr, + ) + sys.exit(1) + + print("\n[example] Captured. Inspect with:") + print(" flowcept-harness sessions") + print(" flowcept-harness show") + print(" flowcept-harness report") + + +if __name__ == "__main__": + main() diff --git a/examples/agents/claude_code/claude_code_example.py b/examples/agents/claude_code/claude_code_example.py new file mode 100644 index 00000000..1965a00d --- /dev/null +++ b/examples/agents/claude_code/claude_code_example.py @@ -0,0 +1,96 @@ +""" +Claude Code provenance capture — install walkthrough and a simulated replay. + +Claude Code delivers each lifecycle event as a JSON object on stdin to a hook +command, one process per event. Capturing a real session needs no code at all: + + /plugin marketplace add # inside Claude Code + /plugin install flowcept + +or, to wire the hooks into settings.json yourself, print them with: + + flowcept-harness install --harness claude_code + +(each hook runs ``flowcept-harness hook --event ``, which reads the +payload from stdin and records it). + +This example replays the hook payloads a short Claude Code session would +deliver — a prompt, a tool call, a subagent, a failed tool, the answer — one +``handle()`` call per event, mimicking the one-process-per-event model. It +runs fully offline: no Claude Code and no API key. Records land as PROV-AGENT +provenance in a JSONL buffer under ``~/.flowcept/harness/buffers/``. + +Run +--- + python examples/agents/claude_code/claude_code_example.py + +Then inspect the capture: + + flowcept-harness sessions + flowcept-harness show + flowcept-harness report + +Plugin configuration +-------------------- +The harness plugins are configured with environment variables, not +settings.yaml (see src/flowcept/agents/harness/README.md for the full table): + + FLOWCEPT_HARNESS_ENABLED=1 # master switch (default) + FLOWCEPT_HARNESS_ONLINE=0 # 1 publishes to a live Flowcept backend + FLOWCEPT_HARNESS_REDACT=1 # redact credential-shaped values + FLOWCEPT_HARNESS_CONTENT=summary # file bodies: full, summary, or none +""" + +from __future__ import annotations + +import uuid + +from flowcept.agents.claude_code.claude_code_plugin import handle +from flowcept.agents.harness.config import load_config + + +def main(): + """Replay one Claude Code session, one hook payload at a time.""" + config = load_config() + session_id = f"claude-code-example-{uuid.uuid4().hex[:8]}" + + def fire(event: str, **fields): + """Deliver one hook payload, shaped exactly as Claude Code sends it.""" + payload = {"hook_event_name": event, "session_id": session_id, "cwd": "/tmp/proj", **fields} + records = handle(payload, config) + print(f"[example] {event:<18} -> {len(records)} record(s)", flush=True) + + # Session opens. + fire("SessionStart", source="startup", model="claude-opus-5") + + # The user asks for something: a turn begins. + fire("UserPromptSubmit", prompt="investigate the flaky network test", prompt_id="p1") + + # The agent greps around — a tool call, Pre and Post in separate processes. + fire("PreToolUse", tool_name="Grep", tool_use_id="t1", tool_input={"pattern": "flaky"}) + fire("PostToolUse", tool_name="Grep", tool_use_id="t1", tool_response={"matches": 3}) + + # It spawns a subagent, whose tool calls land in the subagent's workflow. + fire("PreToolUse", tool_name="Task", tool_use_id="t2", tool_input={"subagent_type": "Explore"}) + fire("SubagentStart", agent_id="a1", agent_type="Explore", task="find the failing test") + fire("PreToolUse", tool_name="Read", tool_use_id="t3", tool_input={"file_path": "tests/test_net.py"}, agent_id="a1") + fire("PostToolUse", tool_name="Read", tool_use_id="t3", tool_response={"lines": 120}, agent_id="a1") + fire("SubagentStop", agent_id="a1", agent_type="Explore", last_assistant_message="tests/test_net.py:42") + fire("PostToolUse", tool_name="Task", tool_use_id="t2", tool_response={"result": "tests/test_net.py:42"}) + + # A command fails: recorded as an errored tool call, not dropped. + fire("PreToolUse", tool_name="Bash", tool_use_id="t4", tool_input={"command": "pytest tests/test_net.py"}) + fire("PostToolUseFailure", tool_name="Bash", tool_use_id="t4", error="exit status 1") + + # The agent answers: the turn closes. Then the session closes. + fire("Stop", last_assistant_message="It is a timing race in tests/test_net.py:42.") + fire("SessionEnd", reason="clear") + + print(f"\n[example] Captured session {session_id!r}. Inspect with:") + print(" flowcept-harness sessions") + print(" flowcept-harness show") + print(" flowcept-harness report") + + +if __name__ == "__main__": + main() diff --git a/examples/agents/cli_harness/cli_harness_example.py b/examples/agents/cli_harness/cli_harness_example.py new file mode 100644 index 00000000..73870345 --- /dev/null +++ b/examples/agents/cli_harness/cli_harness_example.py @@ -0,0 +1,102 @@ +""" +Generic CLI-harness provenance capture through the profile-driven adapter. + +Codex CLI, Gemini CLI, Cursor, OpenCode and friends all have the same shape as +Claude Code — a JSON event handed to a hook command — but disagree on what the +fields are called. The differences live in JSON *profiles* under +``src/flowcept/agents/cli_harness/profiles/``; adding a harness means adding a +file, not writing code. + +In production the harness's hook is pointed at: + + flowcept-harness hook --harness codex --profile codex + +This example is a Python driver that simulates that flow: it replays the hook +payloads a Codex-style session would deliver — one ``handle()`` call per +event, mimicking the one-process-per-event model — using the field names the +bundled ``codex`` profile maps. It runs fully offline, no harness and no API +key. Records land as PROV-AGENT provenance in a JSONL buffer under +``~/.flowcept/harness/buffers/``. + +Run +--- + python examples/agents/cli_harness/cli_harness_example.py + +Then inspect the capture: + + flowcept-harness sessions + flowcept-harness show + flowcept-harness report + +Plugin configuration +-------------------- +The harness plugins are configured with environment variables, not +settings.yaml (see src/flowcept/agents/harness/README.md for the full table): + + FLOWCEPT_HARNESS_ENABLED=1 # master switch (default) + FLOWCEPT_HARNESS_ONLINE=0 # 1 publishes to a live Flowcept backend + FLOWCEPT_HARNESS_REDACT=1 # redact credential-shaped values + +``flowcept-harness install --harness codex`` prints the hook command to wire +into the real harness. +""" + +from __future__ import annotations + +import uuid + +from flowcept.agents.cli_harness.cli_harness_plugin import handle +from flowcept.agents.harness.config import load_config + + +def main(): + """Replay one Codex-style session, one hook payload at a time.""" + config = load_config() + session_id = f"codex-example-{uuid.uuid4().hex[:8]}" + + # Each payload uses Codex's own event and field names; the `codex` profile + # maps them onto the normalized HarnessEvent. The event name is read from + # the payload's "type" key, exactly as a typed notification envelope + # delivers it. + events = [ + # Session opens. + {"type": "session-configured", "session_id": session_id, "cwd": "/tmp/proj", "model": "gpt-5"}, + # The user asks for something: a turn begins. + {"type": "user-message", "session_id": session_id, "message": "run the unit tests"}, + # The agent runs a command: a tool call, begin and end. + { + "type": "exec-command-begin", + "session_id": session_id, + "call_id": "call-1", + "command": "pytest -q tests/unit", + }, + { + "type": "exec-command-end", + "session_id": session_id, + "call_id": "call-1", + "stdout": "42 passed in 3.14s", + }, + # The agent answers: the turn closes. + { + "type": "agent-turn-complete", + "session_id": session_id, + "last_agent_message": "All 42 unit tests pass.", + }, + # Session closes. + {"type": "session-end", "session_id": session_id}, + ] + + total = 0 + for payload in events: + records = handle(payload, config, harness="codex", profile="codex") + total += len(records) + print(f"[example] {payload['type']:<22} -> {len(records)} record(s)", flush=True) + + print(f"\n[example] Captured session {session_id!r} ({total} records). Inspect with:") + print(" flowcept-harness sessions") + print(" flowcept-harness show") + print(" flowcept-harness report") + + +if __name__ == "__main__": + main() diff --git a/examples/agents/combined_agentic_systems/combined_example.py b/examples/agents/combined_agentic_systems/combined_example.py new file mode 100644 index 00000000..30e33999 --- /dev/null +++ b/examples/agents/combined_agentic_systems/combined_example.py @@ -0,0 +1,400 @@ +""" +examples/agents/combined_example.py +===================================== + +Runs all four frameworks simultaneously. Each increments a counter 5 times +(1+2+3+4+5 = 15) concurrently in its own thread. Total = 4 × 15 = 60. +A final llm_chat() call (OpenAI or Anthropic) interprets the combined result. + + Academy ──┐ + AutoGen ──┤ (run in parallel) → combined total = 60 + CrewAI ──┤ + LangGraph ┘ + +Run +--- + OPENAI_API_KEY=sk-... (or ANTHROPIC_API_KEY=sk-ant-...) python examples/agents/combined_example.py + +Plugin configuration +-------------------- +Enable all four plugins in your settings.yaml: + + plugins: + academy: + enabled: true + kind: academy + workflow_name: "combined-academy" + performance_tracking: true + autogen: + enabled: true + kind: autogen + workflow_name: "combined-autogen" + performance_tracking: true + crewai: + enabled: true + kind: crewai + workflow_name: "combined-crewai" + performance_tracking: true + langgraph: + enabled: true + kind: langgraph + workflow_name: "combined-langgraph" + performance_tracking: true + +Flowcept auto-starts/stops all plugins and shares its campaign_id across them. +Access running plugins via flowcept.plugins[""] when needed. +""" +from __future__ import annotations + +import asyncio +import os +import sys +import concurrent.futures +from unittest.mock import patch + +sys.path.insert(0, os.path.dirname(__file__)) +os.environ.setdefault("CREWAI_TRACING_ENABLED", "false") + +from flowcept import Flowcept +# Use whichever LLM provider has a key in the environment. +if os.getenv("ANTHROPIC_API_KEY"): + from flowcept.agents.academy.academy_plugin import anthropic_chat as llm_chat + + LLM_MODEL = "claude-haiku-4-5-20251001" +else: + from flowcept.agents.academy.academy_plugin import openai_chat as llm_chat + + LLM_MODEL = "gpt-4o-mini" + + +# ============================================================ +# ACADEMY +# ============================================================ + +def _run_academy(plugin) -> int: + from academy.agent import Agent, action, loop + from academy.exchange import LocalExchangeFactory + from academy.manager import Manager + from concurrent.futures import ThreadPoolExecutor + + class CounterAgent(Agent): + def __init__(self): + super().__init__() + self._value = 0 + + @action + async def increment(self, n: int = 1) -> int: + self._value += n + return self._value + + @action + async def get_value(self) -> int: + return self._value + + class SummaryAgent(Agent): + def __init__(self): + super().__init__() + self._counter = None + self._result = 0 + self._done = False + + @action + async def set_counter(self, counter) -> None: + self._counter = counter + + @action + async def get_result(self) -> int: + return self._result + + @action + async def is_done(self) -> bool: + return self._done + + @loop + async def run_loop(self, shutdown: asyncio.Event) -> None: + while self._counter is None and not shutdown.is_set(): + await asyncio.sleep(0.05) + if shutdown.is_set(): + return + for step in range(1, 6): + await self._counter.increment(step) + self._result = await self._counter.get_value() + self._done = True + + async def _inner(): + exchange = LocalExchangeFactory() + executor = ThreadPoolExecutor(max_workers=4) + async with await Manager.from_exchange_factory( + factory=exchange, executors=executor + ) as manager: + counter = await manager.launch(CounterAgent) + await counter.ping() + summary = await manager.launch(SummaryAgent) + await summary.ping() + await summary.set_counter(counter) + for _ in range(60): + await asyncio.sleep(0.5) + if await summary.is_done(): + break + return await summary.get_result() + + result = asyncio.run(_inner()) + print(f"[Academy] counter = {result}", flush=True) + return result + + +# ============================================================ +# AUTOGEN +# ============================================================ + +def _run_autogen(plugin) -> int: + try: + from autogen_agentchat.agents import AssistantAgent + from autogen_agentchat.conditions import MaxMessageTermination + from autogen_agentchat.teams import RoundRobinGroupChat + from autogen_core.models import ChatCompletionClient, CreateResult, RequestUsage + except ImportError: + print("[AutoGen] not installed — skipping", file=sys.stderr) + return 0 + + _STEPS = list(range(1, 6)) + _running = 0 + _responses = [] + for s in _STEPS: + _running += s + _responses.append(f"assistant: increment({s}) → total = {_running}") + _responses.append(f"critic: verified. Final = {_running}. TERMINATE") + _idx = [0] + + class _Stub(ChatCompletionClient): + @property + def model_info(self): + from autogen_core.models import ModelInfo + return ModelInfo(vision=False, function_calling=False, json_output=False, + family="stub", structured_output=False) + + @property + def capabilities(self): + from autogen_core.models import ModelCapabilities + return ModelCapabilities(vision=False, function_calling=False, json_output=False) + + async def create(self, *_, **__) -> CreateResult: + content = _responses[_idx[0] % len(_responses)] + _idx[0] += 1 + return CreateResult(content=content, + usage=RequestUsage(prompt_tokens=10, completion_tokens=10), + finish_reason="stop", cached=False) + + async def create_stream(self, *_, **__): + yield await self.create() + + def actual_usage(self): return RequestUsage(0, 0) + def total_usage(self): return RequestUsage(0, 0) + def count_tokens(self, *_, **__): return 0 + def remaining_tokens(self, *_, **__): return 4096 + async def close(self): pass + + client = _Stub() + assistant = AssistantAgent( + name="assistant", model_client=client, + system_message="Increment a counter by the given step and report the total.", + ) + critic = AssistantAgent( + name="critic", model_client=client, + system_message="Verify each total is correct. End with TERMINATE when done.", + ) + team = RoundRobinGroupChat( + participants=[assistant, critic], + termination_condition=MaxMessageTermination(len(_responses) + 1), + ) + + async def _inner(): + result = await plugin.run_team( + team, + "Increment a counter 5 times with step values 1,2,3,4,5 and report totals.", + team_name="counter-team", + ) + total = 0 + for msg in result.messages: + for token in getattr(msg, "content", "").split(): + try: + total = int(token) + except ValueError: + pass + return total + + result = asyncio.run(_inner()) + print(f"[AutoGen] counter = {result}", flush=True) + return result + + +# ============================================================ +# CREWAI +# ============================================================ + +def _run_crewai(_plugin) -> int: + try: + from crewai import Agent, Task, Crew, LLM + except ImportError: + print("[CrewAI] not installed — skipping", file=sys.stderr) + return 0 + + def _fake_completion(content): + from openai.types.chat import ChatCompletion, ChatCompletionMessage + from openai.types.chat.chat_completion import Choice + from openai.types import CompletionUsage + return ChatCompletion( + id="stub", object="chat.completion", created=0, model="gpt-4o-mini", + choices=[Choice(finish_reason="stop", index=0, logprobs=None, + message=ChatCompletionMessage(role="assistant", content=content))], + usage=CompletionUsage(prompt_tokens=20, completion_tokens=20, total_tokens=40), + ) + + _responses = [ + "Final Answer: increment(1)=1 increment(2)=3 increment(3)=6 increment(4)=10 increment(5)=15. Total=15", + "Final Answer: The counter reached 15 = 1+2+3+4+5.", + ] + _idx = [0] + + def _stub(*_, **__): + c = _responses[_idx[0] % len(_responses)] + _idx[0] += 1 + return _fake_completion(c) + + llm = LLM(model="openai/gpt-4o-mini", api_key="stub-key") + counter_agent = Agent( + role="Counter", llm=llm, verbose=False, max_iter=2, + goal="Increment a counter 5 times with steps 1-5 and report each total.", + backstory="You are a precise counter agent.", + ) + summary_agent = Agent( + role="Summariser", llm=llm, verbose=False, max_iter=2, + goal="Summarise the counter result in one sentence.", + backstory="You are a concise analyst.", + ) + count_task = Task( + description="Increment a counter 5 times with step values 1,2,3,4,5 and report each running total.", + expected_output="Step-by-step totals and final value.", + agent=counter_agent, + ) + summary_task = Task( + description="Summarise the counter result in one sentence.", + expected_output="One sentence.", + agent=summary_agent, + context=[count_task], + ) + crew = Crew(agents=[counter_agent, summary_agent], + tasks=[count_task, summary_task], verbose=False) + + with patch("openai.resources.chat.completions.Completions.create", side_effect=_stub): + crew.kickoff() + + result = 15 # 1+2+3+4+5 + print(f"[CrewAI] counter = {result}", flush=True) + return result + + +# ============================================================ +# LANGGRAPH +# ============================================================ + +def _run_langgraph(plugin) -> int: + try: + from langgraph.graph import StateGraph, END + except ImportError: + print("[LangGraph] not installed — skipping", file=sys.stderr) + return 0 + + from typing import TypedDict + + class State(TypedDict): + total: int + steps: list + + def _make_node(step: int): + def node(state: State) -> State: + return {"total": state["total"] + step, "steps": list(state["steps"]) + [step]} + node.__name__ = f"add_{step}" + return node + + builder = StateGraph(State) + for s in range(1, 6): + builder.add_node(f"add_{s}", _make_node(s)) + builder.set_entry_point("add_1") + for s in range(1, 5): + builder.add_edge(f"add_{s}", f"add_{s+1}") + builder.add_edge("add_5", END) + graph = builder.compile() + + final = graph.invoke( + {"total": 0, "steps": []}, + config={"callbacks": [plugin.callback_handler]}, + ) + result = final["total"] + print(f"[LangGraph] counter = {result}", flush=True) + return result + + +# ============================================================ +# Main +# ============================================================ + +def main() -> None: + print("[combined] Starting all four frameworks …\n", flush=True) + + with Flowcept() as flowcept: + academy_plugin = flowcept.plugins.get("academy") + autogen_plugin = flowcept.plugins.get("autogen") + crewai_plugin = flowcept.plugins.get("crewai") + langgraph_plugin = flowcept.plugins.get("langgraph") + + print(f"[combined] campaign_id = {flowcept.campaign_id}\n", flush=True) + + print("[combined] All four frameworks starting concurrently …\n", flush=True) + + results: dict[str, int] = {} + + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool: + futures = { + pool.submit(_run_academy, academy_plugin): "academy", + pool.submit(_run_autogen, autogen_plugin): "autogen", + pool.submit(_run_crewai, crewai_plugin): "crewai", + pool.submit(_run_langgraph, langgraph_plugin): "langgraph", + } + for future in concurrent.futures.as_completed(futures): + name = futures[future] + try: + results[name] = future.result() + except Exception as exc: + print(f"[combined] {name} raised: {exc}", flush=True) + results[name] = 0 + + combined_total = sum(results.values()) + + print("\n[combined] Calling llm_chat() to interpret combined result …", flush=True) + interpretation = llm_chat( + prompt=( + f"Four independent frameworks (Academy, AutoGen, CrewAI, LangGraph) each " + f"incremented a counter 5 times with values 1, 2, 3, 4, 5, each reaching 15. " + f"Combined total = {combined_total}. " + f"In exactly one sentence, explain what this combined result represents." + ), + model=LLM_MODEL, + system="You are a concise data analyst.", + temperature=0.3, + context={"example": "combined", "call_type": "interpret_result"}, + ) + print(f"\n[combined] LLM says: {interpretation}\n", flush=True) + + print("\n" + "=" * 60) + print("COMBINED RESULT") + print("=" * 60) + for name, val in sorted(results.items()): + print(f" {name:<12} {val:>4} (expected 15)") + print(f" {'total':<12} {combined_total:>4} (expected 60 = 4 × 15)") + assert combined_total == 60, f"Expected 60 but got {combined_total}" + print("\n[combined] Assertion passed.", flush=True) + + +if __name__ == "__main__": + main() diff --git a/examples/agents/crewai/crewai_example.py b/examples/agents/crewai/crewai_example.py new file mode 100644 index 00000000..780b58ba --- /dev/null +++ b/examples/agents/crewai/crewai_example.py @@ -0,0 +1,179 @@ +""" +examples/agents/crewai/crewai_example.py +========================================= + +Counter test mirroring the Academy example: + - CounterAgent increments a counter 5 times (values 1-5), reporting each step + - SummaryAgent reads the final total and calls llm_chat() to interpret it, + mirroring SummaryAgent.interpret_result() in the Academy example + +Run +--- + OPENAI_API_KEY=sk-... (or ANTHROPIC_API_KEY=sk-ant-...) python examples/agents/crewai/crewai_example.py + +Plugin configuration +-------------------- +Enable the CrewAI plugin in your settings.yaml: + + plugins: + crewai: + enabled: true + kind: crewai + workflow_name: "crewai-counter-test" + performance_tracking: true + +Flowcept will auto-start/stop the plugin — no explicit plugin.start() / +plugin.stop() calls needed. +""" +from __future__ import annotations + +import os +import sys +from unittest.mock import patch + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) +os.environ.setdefault("CREWAI_TRACING_ENABLED", "false") + +from flowcept import Flowcept +# Use whichever LLM provider has a key in the environment. +if os.getenv("ANTHROPIC_API_KEY"): + from flowcept.agents.crewai.crewai_plugin import anthropic_chat as llm_chat + + LLM_MODEL = "claude-haiku-4-5-20251001" +else: + from flowcept.agents.crewai.crewai_plugin import openai_chat as llm_chat + + LLM_MODEL = "gpt-4o-mini" + +try: + from crewai import Agent, Task, Crew, LLM +except ImportError: + print("ERROR: pip install crewai", file=sys.stderr) + sys.exit(1) + + +# --------------------------------------------------------------------------- +# Stub LLM — canned counter responses, no API key needed for the crew +# --------------------------------------------------------------------------- + +def _make_fake_openai_completion(content: str): + from openai.types.chat import ChatCompletion, ChatCompletionMessage + from openai.types.chat.chat_completion import Choice + from openai.types import CompletionUsage + return ChatCompletion( + id="chatcmpl-stub", + choices=[Choice( + finish_reason="stop", + index=0, + message=ChatCompletionMessage(role="assistant", content=content), + logprobs=None, + )], + created=1_700_000_000, + model="gpt-4o-mini", + object="chat.completion", + usage=CompletionUsage(completion_tokens=30, prompt_tokens=60, total_tokens=90), + ) + + +_STUB_RESPONSES = [ + ( + "Final Answer: Counter incremented 5 times:\n" + " increment(1) → 1\n" + " increment(2) → 3\n" + " increment(3) → 6\n" + " increment(4) → 10\n" + " increment(5) → 15\n" + "Final value: 15" + ), +] +_response_idx = 0 + + +def _stub_create(*args, **kwargs): + global _response_idx + content = _STUB_RESPONSES[_response_idx % len(_STUB_RESPONSES)] + _response_idx += 1 + return _make_fake_openai_completion(content) + + +# --------------------------------------------------------------------------- +# Build crew +# --------------------------------------------------------------------------- + +def build_crew() -> Crew: + llm = LLM(model="openai/gpt-4o-mini", api_key="stub-key") + + counter_agent = Agent( + role="Counter", + goal="Increment a counter 5 times with step values 1, 2, 3, 4, 5 and report each running total.", + backstory=( + "You are a precise counter agent. You add each step value to a running " + "total and report the result after every increment." + ), + llm=llm, + verbose=False, + max_iter=2, + ) + + count_task = Task( + description=( + "Increment a counter 5 times using step values 1, 2, 3, 4, 5. " + "Report the running total after each increment and the final value." + ), + expected_output=( + "A step-by-step report of each increment and the final counter value of 15." + ), + agent=counter_agent, + ) + + return Crew( + agents=[counter_agent], + tasks=[count_task], + verbose=False, + ) + + +# --------------------------------------------------------------------------- +# Main — mirrors academy_example.main() +# --------------------------------------------------------------------------- + +def main() -> None: + crew = build_crew() + + print("\n[example] Running CrewAI crew …\n", flush=True) + + with Flowcept(): + with patch( + "openai.resources.chat.completions.Completions.create", + side_effect=_stub_create, + ): + crew.kickoff() + + final_value = 15 # 1+2+3+4+5 + + # Interpret the result — mirrors SummaryAgent.interpret_result() @action. + # llm_chat() fires record_llm_call() automatically. + print("\n[example] Calling llm_chat() to interpret the counter result …", flush=True) + interpretation = llm_chat( + prompt=( + f"A counter was incremented 5 times with values 1, 2, 3, 4, 5 " + f"and reached a final value of {final_value}. " + f"In exactly one sentence, explain what this arithmetic result represents." + ), + model=LLM_MODEL, + system="You are a concise data analyst.", + temperature=0.3, + context={"agent": "Counter", "call_type": "interpret_result"}, + ) + print(f"\n[example] LLM says: {interpretation}\n", flush=True) + + print("\n" + "=" * 60) + print("CREW OUTPUT") + print("=" * 60) + print(f"\nCounter reached {final_value} (expected 15 = 1+2+3+4+5)", flush=True) + assert final_value == 15, f"Expected 15 but got {final_value}" + print("[example] Assertion passed.", flush=True) + + +if __name__ == "__main__": + main() diff --git a/examples/agents/harness/harness_example.py b/examples/agents/harness/harness_example.py new file mode 100644 index 00000000..7345767e --- /dev/null +++ b/examples/agents/harness/harness_example.py @@ -0,0 +1,49 @@ +"""Example: capture an AI coding-harness session as Flowcept PROV-AGENT provenance. + +The harness plugins (flowcept.agents.harness and the per-harness plugin +modules) capture what an agentic coding session actually did — prompts, turns, +tool calls, subagents — into a JSONL buffer Flowcept reads natively. + +Most captures need no code at all: + +* Claude Code — install the plugin in ``plugins/flowcept`` or run + ``flowcept-harness install`` +* Codex / Gemini / ... — point the harness's hook at + ``flowcept-harness hook --harness codex --profile codex`` +* OpenTelemetry GenAI — ``flowcept.agents.otel.otel_plugin.FlowceptSpanExporter`` +* OpenAI Agents SDK — ``flowcept.agents.openai_agents.openai_agents_plugin.install()`` +* LangChain / LangGraph — ``flowcept.agents.langchain.langchain_plugin.FlowceptCallbackHandler`` +* Claude Agent SDK — ``flowcept.agents.claude_agent_sdk.claude_agent_sdk_plugin.trace_query`` + +This example shows the one case that does need code: an agent you wrote +yourself, driven through :class:`SessionTracer`. Afterwards, inspect the +capture with ``flowcept-harness sessions`` / ``show`` / ``report``. +""" + +from flowcept.agents.harness import SessionTracer + + +def main(): + """Trace a tiny hand-rolled agent session.""" + with SessionTracer("example_agent", model="claude-opus-5") as tracer: + tracer.prompt("summarize the repository") + + with tracer.tool("read_file", {"path": "README.md"}) as call: + call.result({"content": "# Flowcept ..."}) + + with tracer.tool("run_tests", {"suite": "unit"}) as call: + call.result({"passed": 42, "failed": 0}) + + tracer.turn_end( + "The repository is Flowcept; all 42 unit tests pass.", + usage={"input_tokens": 900, "output_tokens": 120}, + ) + + print("Captured. Inspect with:") + print(" flowcept-harness sessions") + print(" flowcept-harness show") + print(" flowcept-harness report") + + +if __name__ == "__main__": + main() diff --git a/examples/agents/langchain/langchain_example.py b/examples/agents/langchain/langchain_example.py new file mode 100644 index 00000000..39e3f2b0 --- /dev/null +++ b/examples/agents/langchain/langchain_example.py @@ -0,0 +1,95 @@ +""" +LangChain / LangGraph provenance capture through FlowceptCallbackHandler. + +A prompt-template -> chat-model chain is invoked twice with the handler passed +as a callback. Each root chain run becomes a turn, and each chat-model run +becomes an ``ai_model_invocation`` task, recorded as PROV-AGENT provenance in a +JSONL buffer under ``~/.flowcept/harness/buffers/``. + +The chat model is ``GenericFakeChatModel`` from langchain-core, so this example +runs fully offline — no API key needed. Swap in ChatOpenAI / ChatAnthropic (or +pass the handler to ``graph.invoke`` for LangGraph) and nothing else changes. + +Run +--- + python examples/agents/langchain/langchain_example.py + +Then inspect the capture: + + flowcept-harness sessions + flowcept-harness show + flowcept-harness report + +Plugin configuration +-------------------- +The harness plugins are configured with environment variables, not +settings.yaml (see src/flowcept/agents/harness/README.md for the full table): + + FLOWCEPT_HARNESS_ENABLED=1 # master switch (default) + FLOWCEPT_HARNESS_ONLINE=0 # 1 publishes to a live Flowcept backend + FLOWCEPT_HARNESS_REDACT=1 # redact credential-shaped values + +No explicit start/stop calls are needed: the handler opens the session when it +is created and closes it when used as a context manager (or via ``close()``). +""" + +from __future__ import annotations + +import sys +import uuid + +from flowcept.agents.langchain.langchain_plugin import FlowceptCallbackHandler + +try: + from langchain_core.language_models.fake_chat_models import GenericFakeChatModel + from langchain_core.messages import AIMessage + from langchain_core.prompts import ChatPromptTemplate +except ImportError: + print("ERROR: pip install langchain-core", file=sys.stderr) + sys.exit(1) + + +def build_chain(): + """Build a prompt-template -> chat-model chain with canned responses.""" + model = GenericFakeChatModel( + messages=iter( + [ + AIMessage(content="A counter incremented 5 times with steps 1..5 reaches 15."), + AIMessage(content="15 is the 5th triangular number: 1+2+3+4+5."), + ] + ) + ) + prompt = ChatPromptTemplate.from_messages( + [ + ("system", "You are a concise data analyst."), + ("human", "{question}"), + ] + ) + return prompt | model + + +def main(): + """Run two turns through the chain, capturing provenance for each.""" + chain = build_chain() + session_id = f"langchain-example-{uuid.uuid4().hex[:8]}" + + # The handler is duck-typed against LangChain's callback protocol: the root + # chain run opens a turn, the chat-model run inside it is recorded as an + # ai_model_invocation, and closing the handler closes the session. + with FlowceptCallbackHandler(session_id=session_id, model="fake-chat-model") as handler: + for question in ( + "A counter was incremented 5 times with values 1, 2, 3, 4, 5. What is the final value?", + "Why is that value 15?", + ): + print(f"\n[example] Question: {question}", flush=True) + answer = chain.invoke({"question": question}, config={"callbacks": [handler]}) + print(f"[example] Answer : {answer.content}", flush=True) + + print(f"\n[example] Captured session {session_id!r}. Inspect with:") + print(" flowcept-harness sessions") + print(" flowcept-harness show") + print(" flowcept-harness report") + + +if __name__ == "__main__": + main() diff --git a/examples/agents/langgraph/langgraph_example.py b/examples/agents/langgraph/langgraph_example.py new file mode 100644 index 00000000..8d3d09d9 --- /dev/null +++ b/examples/agents/langgraph/langgraph_example.py @@ -0,0 +1,162 @@ +""" +examples/agents/langgraph/langgraph_example.py +=============================================== + +Counter test mirroring the Academy example: + + [add_1] → [add_2] → [add_3] → [add_4] → [add_5] → [interpret] + +Each node adds its step value to the running total. +The final node calls llm_chat() to interpret the result (15 = 1+2+3+4+5). + +Run +--- + OPENAI_API_KEY=sk-... (or ANTHROPIC_API_KEY=sk-ant-...) python examples/agents/langgraph/langgraph_example.py + +Plugin configuration +-------------------- +Enable the LangGraph plugin in your settings.yaml: + + plugins: + langgraph: + enabled: true + kind: langgraph + workflow_name: "langgraph-counter-test" + performance_tracking: true + +Flowcept will auto-start/stop the plugin — no explicit plugin.start() / +plugin.stop() calls needed. Access the running plugin via +flowcept.plugins["langgraph"] to retrieve the callback_handler. +""" +from __future__ import annotations + +import os +import sys +from typing import TypedDict + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + +from flowcept import Flowcept +# Use whichever LLM provider has a key in the environment. +if os.getenv("ANTHROPIC_API_KEY"): + from flowcept.agents.langgraph.langgraph_plugin import anthropic_chat as llm_chat + + LLM_MODEL = "claude-haiku-4-5-20251001" +else: + from flowcept.agents.langgraph.langgraph_plugin import openai_chat as llm_chat + + LLM_MODEL = "gpt-4o-mini" + +try: + from langgraph.graph import StateGraph, END +except ImportError: + print( + "ERROR: LangGraph not installed.\n" + " pip install langgraph langchain-core", + file=sys.stderr, + ) + sys.exit(1) + + +# --------------------------------------------------------------------------- +# Graph state +# --------------------------------------------------------------------------- + +class CounterState(TypedDict): + total: int + steps: list[int] + interpretation: str + + +# --------------------------------------------------------------------------- +# Graph nodes — each adds its step value to the running total +# --------------------------------------------------------------------------- + +def _make_add_node(step: int): + def node(state: CounterState) -> CounterState: + new_total = state["total"] + step + steps = list(state["steps"]) + [step] + print(f"[node] increment({step}) → total = {new_total}", flush=True) + return {"total": new_total, "steps": steps} + node.__name__ = f"add_{step}" + return node + + +def interpret_node(state: CounterState) -> CounterState: + """Call the LLM to interpret the final counter value.""" + total = state["total"] + steps = state["steps"] + steps_str = "+".join(str(s) for s in steps) + interpretation = llm_chat( + prompt=( + f"A counter was incremented {len(steps)} times with values {steps_str} " + f"and reached a final value of {total}. " + f"In exactly one sentence, explain what this arithmetic result represents." + ), + model=LLM_MODEL, + system="You are a concise data analyst.", + temperature=0.3, + context={"node": "interpret", "call_type": "interpret_result"}, + ) + print(f"[node] LLM says: {interpretation}", flush=True) + return {"interpretation": interpretation} + + +# --------------------------------------------------------------------------- +# Build the graph +# --------------------------------------------------------------------------- + +def build_graph(): + builder = StateGraph(CounterState) + + for step in range(1, 6): + builder.add_node(f"add_{step}", _make_add_node(step)) + builder.add_node("interpret", interpret_node) + + builder.set_entry_point("add_1") + for step in range(1, 5): + builder.add_edge(f"add_{step}", f"add_{step + 1}") + builder.add_edge("add_5", "interpret") + builder.add_edge("interpret", END) + + return builder.compile() + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main(): + graph = build_graph() + + print("\n[example] Running counter graph …\n", flush=True) + + initial_state: CounterState = { + "total": 0, + "steps": [], + "interpretation": "", + } + + with Flowcept() as flowcept: + plugin = flowcept.plugins.get("langgraph") + if plugin is None: + raise RuntimeError( + "LangGraph plugin is not enabled. Set 'plugins.langgraph.enabled: true' in settings.yaml." + ) + final_state = graph.invoke( + initial_state, + config={"callbacks": [plugin.callback_handler]}, + ) + + print("\n" + "=" * 60) + print("GRAPH OUTPUT") + print("=" * 60) + steps_str = "+".join(str(s) for s in final_state["steps"]) + print(f"\nSteps : {steps_str} = {final_state['total']} (expected 15)") + print(f"\nLLM says : {final_state['interpretation']}") + assert final_state["total"] == 15, f"Expected 15 but got {final_state['total']}" + print("\n[example] Assertion passed.", flush=True) + + +if __name__ == "__main__": + main() diff --git a/examples/agents/openai_agents/openai_agents_example.py b/examples/agents/openai_agents/openai_agents_example.py new file mode 100644 index 00000000..2f31b4ce --- /dev/null +++ b/examples/agents/openai_agents/openai_agents_example.py @@ -0,0 +1,118 @@ +""" +OpenAI Agents SDK provenance capture through FlowceptTraceProcessor. + +The Agents SDK already traces itself — every run produces a trace of typed +spans — and ``install()`` registers Flowcept as one more consumer of those +spans. Nothing about how you call the SDK changes. + +Two modes: + +* ``OPENAI_API_KEY`` set : run a real agent with ``Runner.run_sync`` — the + trace it produces is captured automatically. +* no key (offline) : drive the SDK's own tracing API directly + (``trace`` / ``generation_span`` / + ``function_span``), which exercises exactly the + same capture path without any API call. + +Either way the session lands as PROV-AGENT provenance in a JSONL buffer under +``~/.flowcept/harness/buffers/``. + +Run +--- + python examples/agents/openai_agents/openai_agents_example.py # offline + OPENAI_API_KEY=sk-... python examples/agents/openai_agents/openai_agents_example.py + +Then inspect the capture: + + flowcept-harness sessions + flowcept-harness show + flowcept-harness report + +Plugin configuration +-------------------- +The harness plugins are configured with environment variables, not +settings.yaml (see src/flowcept/agents/harness/README.md for the full table): + + FLOWCEPT_HARNESS_ENABLED=1 # master switch (default) + FLOWCEPT_HARNESS_ONLINE=0 # 1 publishes to a live Flowcept backend + FLOWCEPT_HARNESS_REDACT=1 # redact credential-shaped values + +``install()`` adds Flowcept alongside the SDK's own trace exporter; +``install(replace=True)`` makes Flowcept the only consumer (used in the +offline mode below so the SDK does not try to export to OpenAI). +""" + +from __future__ import annotations + +import os +import sys +import uuid + +from flowcept.agents.openai_agents.openai_agents_plugin import install + +try: + from agents.tracing import function_span, generation_span, trace +except ImportError: + print("ERROR: pip install openai-agents", file=sys.stderr) + sys.exit(1) + + +def run_with_api_key(): + """Run a real agent; its trace is captured by the installed processor.""" + from agents import Agent, Runner + + install() # register once, nothing else changes + + agent = Agent( + name="counter-analyst", + instructions="You are a concise data analyst.", + model="gpt-4o-mini", + ) + result = Runner.run_sync( + agent, + "A counter was incremented 5 times with values 1, 2, 3, 4, 5. " + "In one sentence, what does the final value 15 represent?", + ) + print(f"\n[example] Agent says: {result.final_output}", flush=True) + + +def run_offline(): + """Emit a synthetic trace through the SDK's tracing API — no API call.""" + # replace=True: Flowcept becomes the only trace consumer, so the SDK does + # not also try to export the trace to OpenAI's backend. + install(replace=True) + + session_id = f"openai-agents-example-{uuid.uuid4().hex[:8]}" + + # `group_id` is the SDK's conversation id; the plugin uses it as the + # provenance session id, so a multi-turn conversation is one session. + with trace("counter-run", group_id=session_id): + with function_span("increment_counter", input='{"steps": [1, 2, 3, 4, 5]}') as tool: + tool.span_data.output = '{"total": 15}' + with generation_span( + model="gpt-4o-mini", + input=[{"role": "user", "content": "What does the counter total 15 represent?"}], + output=[{"role": "assistant", "content": "15 is the sum 1+2+3+4+5 — the 5th triangular number."}], + usage={"input_tokens": 25, "output_tokens": 18}, + ): + pass + + print(f"\n[example] Captured synthetic session {session_id!r} (offline mode).", flush=True) + + +def main(): + """Capture one OpenAI Agents SDK session, real or synthetic.""" + if os.getenv("OPENAI_API_KEY"): + run_with_api_key() + else: + print("[example] OPENAI_API_KEY not set — emitting synthetic SDK spans instead.", flush=True) + run_offline() + + print("\n[example] Inspect with:") + print(" flowcept-harness sessions") + print(" flowcept-harness show") + print(" flowcept-harness report") + + +if __name__ == "__main__": + main() diff --git a/examples/agents/otel/otel_example.py b/examples/agents/otel/otel_example.py new file mode 100644 index 00000000..af57217e --- /dev/null +++ b/examples/agents/otel/otel_example.py @@ -0,0 +1,98 @@ +""" +OpenTelemetry GenAI span capture through FlowceptSpanExporter. + +Anything already instrumented with OTel starts producing Flowcept provenance +with three lines of setup: build a ``TracerProvider``, add a +``SimpleSpanProcessor`` wrapping ``FlowceptSpanExporter``, and emit spans. This +example emits synthetic GenAI spans locally — a model call and a tool call — +so it runs fully offline, no collector and no API key. + +Spans are read through the OTel GenAI semantic conventions: +``gen_ai.operation.name`` separates a model call from a tool call, +``gen_ai.tool.name`` names the tool, and ``gen_ai.conversation.id`` groups +spans into a session. Spans without a conversation id, and non-GenAI spans, +are ignored. Records land as PROV-AGENT provenance in a JSONL buffer under +``~/.flowcept/harness/buffers/``. + +To ingest spans a collector already wrote instead (JSON/JSONL, including OTLP +envelopes), use ``flowcept.agents.otel.otel_plugin.ingest_file("spans.jsonl")``. + +Run +--- + python examples/agents/otel/otel_example.py + +Then inspect the capture: + + flowcept-harness sessions + flowcept-harness show + flowcept-harness report + +Plugin configuration +-------------------- +The harness plugins are configured with environment variables, not +settings.yaml (see src/flowcept/agents/harness/README.md for the full table): + + FLOWCEPT_HARNESS_ENABLED=1 # master switch (default) + FLOWCEPT_HARNESS_ONLINE=0 # 1 publishes to a live Flowcept backend + FLOWCEPT_HARNESS_REDACT=1 # redact credential-shaped values +""" + +from __future__ import annotations + +import sys +import uuid + +from flowcept.agents.otel.otel_plugin import FlowceptSpanExporter + +try: + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor +except ImportError: + print("ERROR: pip install opentelemetry-sdk (or: pip install 'flowcept[harness_otel]')", file=sys.stderr) + sys.exit(1) + + +def main(): + """Emit two synthetic GenAI spans through a real OTel tracer provider.""" + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(FlowceptSpanExporter())) + tracer = provider.get_tracer("otel-example") + + session_id = f"otel-example-{uuid.uuid4().hex[:8]}" + + # A tool execution: gen_ai.operation.name "execute_tool" (or a + # gen_ai.tool.name attribute) marks the span as an agent_tool task. + with tracer.start_as_current_span("increment_counter") as span: + span.set_attribute("gen_ai.operation.name", "execute_tool") + span.set_attribute("gen_ai.conversation.id", session_id) + span.set_attribute("gen_ai.tool.name", "increment_counter") + span.set_attribute("gen_ai.tool.call.id", "call-1") + span.set_attribute("gen_ai.tool.call.arguments", '{"steps": [1, 2, 3, 4, 5]}') + span.set_attribute("gen_ai.tool.call.result", '{"total": 15}') + + # A model invocation: gen_ai.operation.name "chat" marks the span as an + # ai_model_invocation task at call granularity. + with tracer.start_as_current_span("chat gpt-4o-mini") as span: + # `gen_ai.system` is optional: the first non-empty value a conversation + # shows is recorded once as the session's provider (first-value-wins). + # Spans with and without it land in the same workflow, so setting it on + # only this span -- as here -- does not split the session. + span.set_attribute("gen_ai.system", "openai") + span.set_attribute("gen_ai.operation.name", "chat") + span.set_attribute("gen_ai.conversation.id", session_id) + span.set_attribute("gen_ai.request.model", "gpt-4o-mini") + span.set_attribute("gen_ai.prompt", "What does the counter total 15 represent?") + span.set_attribute("gen_ai.completion", "15 is the sum 1+2+3+4+5 — the 5th triangular number.") + span.set_attribute("gen_ai.usage.input_tokens", 25) + span.set_attribute("gen_ai.usage.output_tokens", 18) + + provider.shutdown() + + print(f"[example] Captured session {session_id!r} from two GenAI spans. Inspect with:") + print(" flowcept-harness sessions") + print(" flowcept-harness show") + print(" flowcept-harness report") + + +if __name__ == "__main__": + main() diff --git a/examples/agents/prov_analysis/prov_analysis_example.py b/examples/agents/prov_analysis/prov_analysis_example.py new file mode 100644 index 00000000..4db83d42 --- /dev/null +++ b/examples/agents/prov_analysis/prov_analysis_example.py @@ -0,0 +1,146 @@ +""" +Agentic provenance analysis over captured PROV-AGENT records. + +``flowcept.agents.prov_analysis`` turns captured provenance into answers: +what happened, what failed, where the time went, how agents behaved, how +work links across frameworks, and how two runs compare. The same pure +functions back four surfaces — the harness MCP server (``analyze_session``, +``analyze_errors``, ``find_slowest``, ``cross_links``), the Flowcept agent +MCP server (``df_*`` / ``db_*`` analysis tools and ``compare_executions``), +the web chat, and ``flowcept-harness analyze``. + +This example synthesizes two short coding sessions by replaying Claude Code +hook payloads (one ``handle()`` call per event, exactly as the hooks deliver +them), plants one framework-plugin record that links back to a harness tool +task, then runs every ``prov_analysis.core`` function over the captured +records and prints the results. It runs fully offline: no Claude Code, no +API key, no database. Records also land as PROV-AGENT provenance in JSONL +buffers under ``~/.flowcept/harness/buffers/``. + +Run +--- + python examples/agents/prov_analysis/prov_analysis_example.py + +Then inspect the same captures from the CLI: + + flowcept-harness sessions + flowcept-harness analyze # summary of the most recent session + flowcept-harness analyze --errors # failure clustering + flowcept-harness analyze --slowest 5 + flowcept-harness analyze --links + +Plugin configuration +-------------------- +The harness plugins are configured with environment variables, not +settings.yaml (see src/flowcept/agents/harness/README.md for the full table): + + FLOWCEPT_HARNESS_ENABLED=1 # master switch (default) + FLOWCEPT_HARNESS_ONLINE=0 # 1 publishes to a live Flowcept backend + FLOWCEPT_HARNESS_REDACT=1 # redact credential-shaped values +""" + +from __future__ import annotations + +import json +import uuid + +from flowcept.agents.claude_code.claude_code_plugin import handle +from flowcept.agents.harness.config import load_config +from flowcept.agents.prov_analysis.core import ( + analyze_agent_behavior, + analyze_errors, + compare_executions, + cross_framework_links, + find_slowest_tasks, + load_records, + summarize_execution, +) + + +def record_session(config, session_id: str, flaky: bool) -> list[dict]: + """Replay one short Claude Code session and return its captured records. + + The ``flaky`` variant fails its second tool call, so the two sessions + differ in error rate — which ``compare_executions`` then surfaces. + """ + records: list[dict] = [] + + def fire(event: str, **fields): + payload = {"hook_event_name": event, "session_id": session_id, "cwd": "/tmp/proj", **fields} + records.extend(handle(payload, config)) + + fire("SessionStart", source="startup", model="claude-opus-5") + fire("UserPromptSubmit", prompt="fix the flaky network test", prompt_id="p1") + fire("PreToolUse", tool_name="Bash", tool_use_id="t1", tool_input={"command": "pytest -q"}) + fire("PostToolUse", tool_name="Bash", tool_use_id="t1", tool_response={"exit_code": 0}) + fire("PreToolUse", tool_name="Bash", tool_use_id="t2", tool_input={"command": "ruff check"}) + if flaky: + fire("PostToolUseFailure", tool_name="Bash", tool_use_id="t2", error="1 error found") + else: + fire("PostToolUse", tool_name="Bash", tool_use_id="t2", tool_response={"exit_code": 0}) + fire("SubagentStart", agent_id="a1", agent_type="Explore") + fire("PreToolUse", tool_name="Grep", tool_use_id="t3", tool_input={"pattern": "flaky"}, agent_id="a1") + fire("PostToolUse", tool_name="Grep", tool_use_id="t3", tool_response={"matches": 2}, agent_id="a1") + fire("SubagentStop", agent_id="a1", agent_type="Explore") + fire("Stop", last_assistant_message="Fixed the race.") + fire("SessionEnd", reason="clear") + return records + + +def plant_cross_framework_record(records: list[dict]) -> None: + """Append a framework-plugin record linking back to a harness tool task. + + This mimics what the LangGraph plugin writes when a graph run is started + by another framework's agent: the parent task id travels in + ``custom_metadata.source_agent_id``, giving ``cross_framework_links`` an + explicit edge to walk. + """ + tool_task = next(r for r in records if r.get("subtype") == "agent_tool") + records.append( + { + "task_id": "lg-1", + "workflow_id": "wf-langgraph", + "activity_id": "graph", + "subtype": "langgraph_graph", + "status": "FINISHED", + "custom_metadata": {"source_agent_id": tool_task["task_id"]}, + } + ) + + +def show(title: str, payload) -> None: + """Print one analysis result as indented JSON under a heading.""" + print(f"\n=== {title} ===") + print(json.dumps(payload, indent=2, default=str)) + + +def main(): + """Capture two synthetic sessions and run every analysis over them.""" + config = load_config() + run_id = uuid.uuid4().hex[:8] + + # Capture two sessions: one clean, one with a failing tool call. + records_clean = record_session(config, f"prov-analysis-example-a-{run_id}", flaky=False) + records_flaky = record_session(config, f"prov-analysis-example-b-{run_id}", flaky=True) + plant_cross_framework_record(records_flaky) + + # load_records also accepts a JSONL buffer path (jsonl_path=...) or a + # workflow/campaign id to pull from the Flowcept DB; here the records + # captured above are passed straight through. + records = load_records(records=records_flaky) + + show("summarize_execution", summarize_execution(records)) + show("analyze_errors", analyze_errors(records)) + show("analyze_agent_behavior", analyze_agent_behavior(records)) + show("find_slowest_tasks (top 3)", find_slowest_tasks(records, limit=3)) + show("cross_framework_links", cross_framework_links(records)) + show("compare_executions (clean vs flaky)", compare_executions(records_clean, records_flaky)) + + print("\n[example] Captured two sessions. Inspect them from the CLI with:") + print(" flowcept-harness sessions") + print(" flowcept-harness analyze") + print(" flowcept-harness analyze --errors") + + +if __name__ == "__main__": + main() diff --git a/plugins/flowcept/.claude-plugin/plugin.json b/plugins/flowcept/.claude-plugin/plugin.json new file mode 100644 index 00000000..c5e3341d --- /dev/null +++ b/plugins/flowcept/.claude-plugin/plugin.json @@ -0,0 +1,18 @@ +{ + "name": "flowcept", + "description": "Capture Claude Code sessions as Flowcept PROV-AGENT provenance and analyze them in place: prompts, turns, tool calls, and subagents written to a JSONL buffer Flowcept reads natively, plus MCP provenance tools (flowcept-provenance server), analysis skills, and an optional auto-report on session end.", + "version": "0.2.0", + "author": { + "name": "flowcept-harness" + }, + "homepage": "https://flowcept.org/", + "repository": "https://github.com/ORNL/flowcept", + "license": "MIT", + "keywords": [ + "provenance", + "flowcept", + "prov-agent", + "observability", + "telemetry" + ] +} diff --git a/plugins/flowcept/.mcp.json b/plugins/flowcept/.mcp.json new file mode 100644 index 00000000..3aac53b6 --- /dev/null +++ b/plugins/flowcept/.mcp.json @@ -0,0 +1,8 @@ +{ + "mcpServers": { + "flowcept-provenance": { + "command": "${CLAUDE_PLUGIN_ROOT}/scripts/mcp-server.sh", + "args": [] + } + } +} diff --git a/plugins/flowcept/hooks/hooks.json b/plugins/flowcept/hooks/hooks.json new file mode 100644 index 00000000..86f7ef1f --- /dev/null +++ b/plugins/flowcept/hooks/hooks.json @@ -0,0 +1,159 @@ +{ + "hooks": { + "SessionStart": [ + { + "hooks": [ + { + "type": "command", + "command": "\"${CLAUDE_PLUGIN_ROOT}/scripts/hook.sh\" SessionStart", + "timeout": 5 + } + ] + } + ], + "SessionEnd": [ + { + "hooks": [ + { + "type": "command", + "command": "\"${CLAUDE_PLUGIN_ROOT}/scripts/hook.sh\" SessionEnd", + "timeout": 10 + } + ] + }, + { + "hooks": [ + { + "type": "command", + "command": "\"${CLAUDE_PLUGIN_ROOT}/scripts/autoreport.sh\"", + "timeout": 30 + } + ] + } + ], + "UserPromptSubmit": [ + { + "hooks": [ + { + "type": "command", + "command": "\"${CLAUDE_PLUGIN_ROOT}/scripts/hook.sh\" UserPromptSubmit", + "timeout": 5 + } + ] + } + ], + "Stop": [ + { + "hooks": [ + { + "type": "command", + "command": "\"${CLAUDE_PLUGIN_ROOT}/scripts/hook.sh\" Stop", + "timeout": 5 + } + ] + } + ], + "StopFailure": [ + { + "hooks": [ + { + "type": "command", + "command": "\"${CLAUDE_PLUGIN_ROOT}/scripts/hook.sh\" StopFailure", + "timeout": 5 + } + ] + } + ], + "PreToolUse": [ + { + "matcher": "*", + "hooks": [ + { + "type": "command", + "command": "\"${CLAUDE_PLUGIN_ROOT}/scripts/hook.sh\" PreToolUse", + "timeout": 5 + } + ] + } + ], + "PostToolUse": [ + { + "matcher": "*", + "hooks": [ + { + "type": "command", + "command": "\"${CLAUDE_PLUGIN_ROOT}/scripts/hook.sh\" PostToolUse", + "timeout": 5 + } + ] + } + ], + "PostToolUseFailure": [ + { + "matcher": "*", + "hooks": [ + { + "type": "command", + "command": "\"${CLAUDE_PLUGIN_ROOT}/scripts/hook.sh\" PostToolUseFailure", + "timeout": 5 + } + ] + } + ], + "SubagentStart": [ + { + "hooks": [ + { + "type": "command", + "command": "\"${CLAUDE_PLUGIN_ROOT}/scripts/hook.sh\" SubagentStart", + "timeout": 5 + } + ] + } + ], + "SubagentStop": [ + { + "hooks": [ + { + "type": "command", + "command": "\"${CLAUDE_PLUGIN_ROOT}/scripts/hook.sh\" SubagentStop", + "timeout": 5 + } + ] + } + ], + "Notification": [ + { + "hooks": [ + { + "type": "command", + "command": "\"${CLAUDE_PLUGIN_ROOT}/scripts/hook.sh\" Notification", + "timeout": 5 + } + ] + } + ], + "PreCompact": [ + { + "hooks": [ + { + "type": "command", + "command": "\"${CLAUDE_PLUGIN_ROOT}/scripts/hook.sh\" PreCompact", + "timeout": 5 + } + ] + } + ], + "PostCompact": [ + { + "hooks": [ + { + "type": "command", + "command": "\"${CLAUDE_PLUGIN_ROOT}/scripts/hook.sh\" PostCompact", + "timeout": 5 + } + ] + } + ] + } +} diff --git a/plugins/flowcept/scripts/autoreport.sh b/plugins/flowcept/scripts/autoreport.sh new file mode 100755 index 00000000..7c7380ab --- /dev/null +++ b/plugins/flowcept/scripts/autoreport.sh @@ -0,0 +1,96 @@ +#!/usr/bin/env bash +# Optional auto-report on SessionEnd. Off by default: does nothing unless +# FLOWCEPT_HARNESS_AUTOREPORT=1. When enabled, generates a Flowcept workflow +# card for the session that just ended and writes it to +# $FLOWCEPT_HARNESS_HOME/reports/.md. +# +# Same discipline as hook.sh: +# +# 1. Never fail. Every path ends in `exit 0`. +# 2. Never write to stdout. Hook stdout can be injected into model context. +# 3. Never block the user. This runs on SessionEnd only, after the session +# is over, so a report is allowed to take seconds -- but still capped by +# the hook timeout in hooks.json. + +set -u + +# Off by default: consume stdin (some CLIs treat an unread pipe as an error) +# and leave silently. +if [ "${FLOWCEPT_HARNESS_AUTOREPORT:-0}" != "1" ]; then + cat >/dev/null 2>&1 || true + exit 0 +fi + +# --- locate an interpreter (same resolution as hook.sh) ---------------------- +PY="${FLOWCEPT_HARNESS_PYTHON:-}" +if [ -z "$PY" ]; then + for candidate in python3 python; do + if command -v "$candidate" >/dev/null 2>&1; then + PY="$candidate" + break + fi + done +fi +[ -z "$PY" ] && exit 0 + +# --- locate the package (same resolution as hook.sh) ------------------------- +ROOT="${CLAUDE_PLUGIN_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}" +for path in "$ROOT/vendor" "$ROOT/../../src"; do + if [ -d "$path/flowcept" ]; then + PYTHONPATH="${PYTHONPATH:+$PYTHONPATH:}$path" + export PYTHONPATH + break + fi +done + +# Map the SessionEnd payload's session_id to its buffer, then run the report +# through `flowcept-harness report --input -o .md>`. +# Falls back to the most recent buffer when the payload is unusable. stdout to +# /dev/null enforces rule 2; failures land in the log, never on the user. +"$PY" - <<'PYEOF' >/dev/null 2>>"${FLOWCEPT_HARNESS_LOG:-/dev/null}" || true +import json +import os +import pathlib +import sys + +home = pathlib.Path( + os.environ.get("FLOWCEPT_HARNESS_HOME", pathlib.Path.home() / ".flowcept" / "harness") +).expanduser() + +try: + payload = json.load(sys.stdin) +except Exception: + payload = {} +session_id = payload.get("session_id") if isinstance(payload, dict) else None + +buffers = sorted(home.glob("buffers/*.jsonl"), key=lambda p: p.stat().st_mtime, reverse=True) +buf = None +if session_id: + for path in buffers: + try: + for line in path.read_text(encoding="utf-8").splitlines(): + line = line.strip() + if not line: + continue + record = json.loads(line) + if record.get("type") == "workflow" and (record.get("used") or {}).get("session_id") == session_id: + buf = path + break + except Exception: + continue + if buf is not None: + break +if buf is None and buffers: + buf = buffers[0] +if buf is None: + raise SystemExit(0) + +reports = home / "reports" +reports.mkdir(parents=True, exist_ok=True) +out = reports / (buf.stem + ".md") + +from flowcept.agents.harness.cli import main + +raise SystemExit(main(["report", "--input", str(buf), "-o", str(out)])) +PYEOF +exit 0 diff --git a/plugins/flowcept/scripts/hook.sh b/plugins/flowcept/scripts/hook.sh new file mode 100755 index 00000000..48dc8f3c --- /dev/null +++ b/plugins/flowcept/scripts/hook.sh @@ -0,0 +1,47 @@ +#!/usr/bin/env bash +# Provenance hook shim. Reads a Claude Code hook payload on stdin and records +# it. Argument 1 is the hook event name, used as a fallback on CLI versions +# that omit `hook_event_name` from the payload. +# +# Three rules this script exists to enforce: +# +# 1. Never fail. A non-zero exit from a hook is surfaced to the user, and on +# some events a hook can block the turn outright. Capture is not worth +# interrupting anyone's work, so every path ends in `exit 0`. +# 2. Never write to stdout. On UserPromptSubmit and SessionStart, a hook's +# stdout is injected into the model's context. +# 3. Never be slow. The capture path is stdlib-only precisely so this can be +# a bare interpreter start with no third-party imports. + +set -u + +# --- locate an interpreter -------------------------------------------------- +# FLOWCEPT_HARNESS_PYTHON wins, for pyenv/conda setups where `python3` on PATH +# is not the one with flowcept-harness installed. +PY="${FLOWCEPT_HARNESS_PYTHON:-}" +if [ -z "$PY" ]; then + for candidate in python3 python; do + if command -v "$candidate" >/dev/null 2>&1; then + PY="$candidate" + break + fi + done +fi +[ -z "$PY" ] && exit 0 + +# --- locate the package ----------------------------------------------------- +# Prefer an installed flowcept; fall back to the repo checkout the plugin +# ships in, so the plugin works with nothing pip-installed at all. +ROOT="${CLAUDE_PLUGIN_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}" +for path in "$ROOT/vendor" "$ROOT/../../src"; do + if [ -d "$path/flowcept" ]; then + PYTHONPATH="${PYTHONPATH:+$PYTHONPATH:}$path" + export PYTHONPATH + break + fi +done + +# stdout to /dev/null enforces rule 2 even if something downstream prints. +"$PY" -m flowcept.agents.claude_code.claude_code_plugin --event "${1:-}" \ + >/dev/null 2>>"${FLOWCEPT_HARNESS_LOG:-/dev/null}" || true +exit 0 diff --git a/plugins/flowcept/scripts/mcp-server.sh b/plugins/flowcept/scripts/mcp-server.sh new file mode 100755 index 00000000..d30ea222 --- /dev/null +++ b/plugins/flowcept/scripts/mcp-server.sh @@ -0,0 +1,40 @@ +#!/usr/bin/env bash +# Launch the flowcept-provenance MCP server over stdio. +# +# Interpreter and package resolution mirror hook.sh: FLOWCEPT_HARNESS_PYTHON +# wins, then python3/python on PATH; an installed flowcept wins, then the repo +# checkout this plugin ships in. Unlike hook.sh, stdout is NOT suppressed -- +# stdout *is* the MCP stdio transport. Diagnostics go to stderr, which the MCP +# client logs. +# +# The module runs as __main__ (it has an `if __name__ == "__main__"` guard), so +# `python -m flowcept.agents.harness.mcp_server` is equivalent to the console +# script `flowcept-harness-mcp`; the module form works without an entry-point +# install. It needs the `mcp` package: pip install 'flowcept[dev]' or `mcp`. + +set -u + +PY="${FLOWCEPT_HARNESS_PYTHON:-}" +if [ -z "$PY" ]; then + for candidate in python3 python; do + if command -v "$candidate" >/dev/null 2>&1; then + PY="$candidate" + break + fi + done +fi +if [ -z "$PY" ]; then + echo "flowcept-provenance: no python interpreter found" >&2 + exit 1 +fi + +ROOT="${CLAUDE_PLUGIN_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}" +for path in "$ROOT/vendor" "$ROOT/../../src"; do + if [ -d "$path/flowcept" ]; then + PYTHONPATH="${PYTHONPATH:+$PYTHONPATH:}$path" + export PYTHONPATH + break + fi +done + +exec "$PY" -m flowcept.agents.harness.mcp_server --transport stdio diff --git a/pyproject.toml b/pyproject.toml index f763ce57..cc9c4706 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,6 +62,12 @@ report_pdf = ["matplotlib", "reportlab", "networkx"] mongo = ["pymongo", "pyarrow"] dask = ["tomli", "dask[distributed]<=2024.10.0"] docs = ["sphinx", "furo"] +diaspora = ["diaspora-stream-api @ git+https://github.com/diaspora-project/diaspora-stream-api"] +# For the AI-coding-harness plugins (flowcept.agents.harness and friends). The +# capture path is deliberately dependency-free; these extras are only for the +# OTel span exporter and the Claude Agent SDK `trace_query` wrapper. +harness_otel = ["opentelemetry-sdk>=1.20.0"] +harness_claude_sdk = ["claude-agent-sdk>=0.1.0"] kafka = ["confluent-kafka<=2.8.0"] # As of today, 2/28/2025, version 2.8.1 is stale. When this gets fixed, let's remove the version constraint. https://pypi.org/project/confluent-kafka/#history rabbitmq = ["pika"] mlflow = ["mlflow-skinny", "SQLAlchemy", "alembic", "watchdog", "cryptography"] @@ -120,6 +126,11 @@ all = [ [tool.hatch.version] path = "src/flowcept/version.py" + +[tool.hatch.metadata] +allow-direct-references = true + + [tool.ruff] line-length = 120 @@ -150,3 +161,9 @@ filterwarnings = [ [project.scripts] flowcept = "flowcept.cli:main" +# AI coding harness provenance capture (flowcept.agents.harness). The capture +# path is stdlib-only, so these need no extras; the MCP server needs [llm_agent] +# (for mcp) and the OTel exporter needs [harness_otel]. +flowcept-harness = "flowcept.agents.harness.cli:main" +flowcept-harness-mcp = "flowcept.agents.harness.mcp_server:main" +flowcept-claude-code = "flowcept.agents.claude_code.claude_code_plugin:main" diff --git a/resources/diaspora/consumer.py b/resources/diaspora/consumer.py new file mode 100644 index 00000000..198d56c1 --- /dev/null +++ b/resources/diaspora/consumer.py @@ -0,0 +1,43 @@ +import json +import os +import time +import csv +from diaspora_stream.api import Driver +print("about to start", flush=True) + +driver_options = { + "root_path": "/tmp/diaspora-data/", +} +driver = Driver(backend="files", options=driver_options) +# create a topic +topic_name = "interception" +consumer_name = "flowcept" +topic = driver.open_topic(topic_name) + +consumer = topic.consumer(name=consumer_name) + +# Get the list of files in the current directory +csv_files = [file for file in os.listdir() if file.endswith('.csv')] +threshold = len(csv_files) +print("about to start with breakpoint ",threshold, flush=True) +while True: + data = [] + metadata = [] + t1 = time.time() + f = consumer.pull() + event = f.wait(timeout_ms=1) + while not f.completed(): + event = f.wait(timeout_ms=10) + t2 = time.time() + if event: + e = event.metadata + print(e) + else: + print("END event") + break + + + + + + diff --git a/resources/diaspora/diaspora_setup.sh b/resources/diaspora/diaspora_setup.sh new file mode 100644 index 00000000..bdfe6b3e --- /dev/null +++ b/resources/diaspora/diaspora_setup.sh @@ -0,0 +1,9 @@ +diaspora-ctl topic create --name interception \ + --driver files \ + --driver.root_path /tmp/diaspora-data/ \ + --topic.num_partitions 1 +sleep 1 + +echo "Created topic." +while true; do sleep 3600; done + diff --git a/resources/sample_settings.yaml b/resources/sample_settings.yaml index a479141e..28b90022 100644 --- a/resources/sample_settings.yaml +++ b/resources/sample_settings.yaml @@ -1,4 +1,4 @@ -flowcept_version: 0.10.8 # Version of the Flowcept package. Do not update this manually. The CI updates it. This setting file is compatible with this version. +flowcept_version: 1.0.4 # Version of the Flowcept package. Do not update this manually. The CI updates it. This setting file is compatible with this version. project: debug: true # Toggle debug mode. This will add a property `debug: true` to all saved data, making it easier to retrieve/delete them later. @@ -44,12 +44,11 @@ experiment: mq: enabled: false - type: redis # or kafka, mofka, rabbitmq; adjust port accordingly (redis: 6379, kafka: 9092, rabbitmq: 5672). If mofka, also set group_file. + type: redis # or kafka, mofka, rabbitmq, diaspora; adjust port accordingly (redis: 6379, kafka: 9092, rabbitmq: 5672). If mofka, also set group_file. host: localhost # uri: ? # instances: ["localhost:6379"] # We can have multiple MQ instances being accessed by the consumers but each interceptor will currently access one single MQ.. port: 6379 - # group_id: auto # Kafka-only consumer group id. Use "auto" to generate a unique group per run. # group_file: mofka.json # username: guest # RabbitMQ only (AMQP); default is "guest" # vhost: / # RabbitMQ only; default is "/" @@ -112,7 +111,6 @@ agent: agent_mode: disabled # How the MCP agent is deployed: disabled | separate | colocated databases: - lmdb: enabled: false path: flowcept_lmdb @@ -129,6 +127,46 @@ databases: # lock_file_path: /var/run/mongod.pid +plugins: + # Agent-framework plugins — set enabled: true to activate without any code changes. + # When Flowcept starts (via context manager or .start()), enabled plugins are + # automatically started and stopped alongside Flowcept. No need to call + # plugin.start() / plugin.stop() or pass a config dict in your code. + # + # Each entry key is an arbitrary name; "kind" selects the plugin class: + # academy → FlowceptAcademyPlugin + # langgraph → FlowceptLangGraphPlugin + # crewai → FlowceptCrewAIPlugin + # autogen → FlowceptAutoGenPlugin + # + # All other keys are forwarded as the plugin's config dict. + + academy: + enabled: false + kind: academy + workflow_name: "academy-workflow" + performance_tracking: true + # perf_csv: "provenance_perf.csv" + # campaign_id: ~ + + # langgraph: + # enabled: true + # kind: langgraph + # workflow_name: "langgraph-workflow" + # performance_tracking: true + + # crewai: + # enabled: true + # kind: crewai + # workflow_name: "crewai-workflow" + # performance_tracking: true + + # autogen: + # enabled: true + # kind: autogen + # workflow_name: "autogen-workflow" + # performance_tracking: true + adapters: # For each key below, you can have multiple instances. Like mlflow1, mlflow2; zambeze1, zambeze2. Use an empty dict, {}, if you won't use any adapter. diff --git a/src/flowcept/agents/README.md b/src/flowcept/agents/README.md index 36fa1fbf..0f599179 100644 --- a/src/flowcept/agents/README.md +++ b/src/flowcept/agents/README.md @@ -1,13 +1,25 @@ -# Flowcept Agent +# Flowcept Agents -This package contains the Flowcept MCP server, client helpers, data-query tools, -MCP-wrapper tools, prompts, context manager, and LLM infrastructure. +This package contains everything agent-related in Flowcept: -For code-assistant behavior, use the repository root `AGENTS.md`. Runtime usage -docs live in `docs/agent.rst`. +1. **The Flowcept Agent** — the MCP server, client helpers, data-query tools, + MCP-wrapper tools, prompts, context manager, and LLM infrastructure that let + an LLM query provenance data. Runtime usage docs live in `docs/agent.rst`. +2. **Agentic framework provenance plugins** — zero-code-change capture of + agent/action/LLM-call provenance from Academy, LangGraph, CrewAI, and + AutoGen. Docs: `docs/agent_plugins.rst`. +3. **AI coding harness provenance plugins** — PROV-AGENT capture of coding + sessions from Claude Code, other CLI harnesses (Codex, Gemini, Cursor, + OpenCode), the Claude Agent SDK, the OpenAI Agents SDK, + LangChain/LangGraph callbacks, and OpenTelemetry GenAI spans. Docs: + `docs/harness_plugins.rst` and [`harness/README.md`](harness/README.md). + +For code-assistant behavior, use the repository root `AGENTS.md`. ## What Lives Here +### Flowcept Agent (MCP server + web chat) + - `chat_orchestration/`: LangChain / LangGraph orchestration for the web chat. This is where the chat runtime, tool routing, and turn-level orchestration live. It should stay separate from HTTP route handlers. @@ -18,6 +30,11 @@ docs live in `docs/agent.rst`. - `data_query_tools/`: shared query logic. This is where task, workflow, object, and DataFrame query behavior lives. These modules can call `DBAPI` for persisted data or read the in-memory DataFrame / workflow object for runtime questions. +- `prov_analysis/`: agentic provenance analysis over PROV-AGENT records — pure + analysis functions (summaries, errors, agent behavior, slowest tasks, + cross-framework links, run comparison) shared by the harness MCP server, the + Flowcept agent MCP server, the web chat, and `flowcept-harness analyze`. + Full docs: [`prov_analysis/README.md`](prov_analysis/README.md). - `prompts/`: prompt-builder functions and prompt registrations. Keep them as plain Python builders that return strings, not Jinja templates. - `provenance_schema_manager/`: schema introspection and documentation context used by @@ -25,6 +42,56 @@ docs live in `docs/agent.rst`. - `llm/`: model construction and normalization helpers. Centralize LLM creation here. - `gui/`: legacy UI helpers. Do not extend this unless the old GUI is being revived. +### Agentic framework provenance plugins + +Each plugin captures intra-agent, inter-agent, and LLM-call provenance from a +framework with zero code changes. They can be auto-started from the `plugins:` +block in `settings.yaml` (kinds: `academy`, `langgraph`, `crewai`, `autogen`) +or constructed directly; all four can share one `campaign_id` via +`from_academy_plugin()`. Each module also exports `openai_chat`, +`anthropic_chat`, and `FlowceptAnthropicClient` LLM-call wrappers. + +- `academy/`: `FlowceptAcademyPlugin` — wraps Academy agents; records + `academy_action` / `academy_loop` / `academy_lifecycle` tasks and nested + `llm_call` records. +- `autogen/`: `FlowceptAutoGenPlugin`, `run_team()`, `FlowceptModelClient` — + records `autogen_run` → `autogen_message` → `llm_call`. +- `crewai/`: `FlowceptCrewAIPlugin` — records `crewai_crew`, `crewai_task` → + `crewai_agent` → `llm_call` / `tool_call` via CrewAI listeners/hooks. +- `langgraph/`: `FlowceptLangGraphPlugin` — records `langgraph_graph` → + `langgraph_node` → `llm_call` / `tool_call` through `plugin.callback_handler`. + +Docs: `docs/agent_plugins.rst`. Runnable examples: `examples/agents/` +(`academy/`, `langgraph/`, `crewai/`, `autogen/`, +`combined_agentic_systems/`). + +### AI coding harness provenance plugins + +Shared capture core plus one adapter module per source. The capture path is +stdlib-only and buffers JSONL per session under `~/.flowcept/harness/buffers/`; +inspect with the `flowcept-harness` CLI (`sessions`, `show`, `status`, +`report`, `flush`, `repair`, `install`, `hook`). + +- `harness/`: the shared capture core — events, recorder, PROV-AGENT record + constructors, JSONL buffers, `SessionTracer` for agents you write yourself, + the `flowcept-harness` CLI (`harness/cli.py`), and the + `flowcept-harness-mcp` server. Full docs: [`harness/README.md`](harness/README.md). +- `claude_code/`: Claude Code hook adapter (used by the `plugins/flowcept` + Claude Code plugin, or by hooks in `settings.json`). +- `cli_harness/`: profile-driven adapter for other CLI harnesses; + `profiles/` ships `codex`, `gemini`, `cursor`, `opencode` JSON profiles. +- `claude_agent_sdk/`: `trace_query` — a drop-in for `claude_agent_sdk.query` + (`pip install "flowcept[harness_claude_sdk]"`). +- `openai_agents/`: `FlowceptTraceProcessor` + `install()` — a tracing + processor for the OpenAI Agents SDK. +- `langchain/`: `FlowceptCallbackHandler` — a LangChain / LangGraph callback + handler that records turns, model calls, and tool calls. +- `otel/`: `FlowceptSpanExporter` and `ingest_file()` for OpenTelemetry GenAI + spans (`pip install "flowcept[harness_otel]"`). + +Docs: `docs/harness_plugins.rst`. Example: +`examples/agents/harness/harness_example.py`. + ## Directory Layout ``` @@ -51,6 +118,10 @@ agents/ df_query_tools.py # DFQueryTools + run_df_query, execute_df_code, generate_result_df, … pandas_utils.py # safe_execute, normalize_output, format_result_df, … + prov_analysis/ # Pure provenance-analysis functions shared by all surfaces + core.py # summarize_execution, analyze_errors, find_slowest_tasks, … + tools.py # ToolResult wrappers over core (no MCP/LangChain imports) + mcp/ mcp_server.py # MCP server entry point (start with `flowcept --start --agent`) mcp_client.py # Client helpers: run_tool() @@ -68,6 +139,21 @@ agents/ db_query_prompts.py # build_db_schema_context df_query_prompts.py # build_pandas_code_prompt, build_plot_code_prompt, … chat_prompts.py # build_chat_system_prompt() for the webservice chat + + # Agentic framework provenance plugins (docs/agent_plugins.rst) + academy/ # FlowceptAcademyPlugin + autogen/ # FlowceptAutoGenPlugin, run_team() + crewai/ # FlowceptCrewAIPlugin + langgraph/ # FlowceptLangGraphPlugin + + # AI coding harness provenance plugins (docs/harness_plugins.rst) + harness/ # shared capture core, SessionTracer, flowcept-harness CLI, MCP server + claude_code/ # Claude Code hook adapter + cli_harness/ # profile-driven adapter (+ profiles/: codex, gemini, cursor, opencode) + claude_agent_sdk/ # trace_query wrapper + openai_agents/ # OpenAI Agents SDK tracing processor + langchain/ # LangChain / LangGraph callback handler + otel/ # OTel GenAI span exporter and ingest ``` ## One Agent, Two Orchestrators @@ -131,6 +217,7 @@ code needed; tool calls are stored automatically when the interceptor is running ```python from flowcept.instrumentation.flowcept_agent_task import FlowceptLLM + wrapped = FlowceptLLM(llm, agent_id=my_agent_id) response = wrapped.invoke("How many tasks failed?") ``` diff --git a/src/flowcept/agents/__init__.py b/src/flowcept/agents/__init__.py index b248dd24..59418367 100644 --- a/src/flowcept/agents/__init__.py +++ b/src/flowcept/agents/__init__.py @@ -1,7 +1,29 @@ # flake8: noqa: F403 -"""Agents subpackage.""" +"""Agents subpackage. -from flowcept.agents.tool_result import ToolResult # noqa: F401 -from flowcept.agents.mcp.mcp_tools import * -from flowcept.agents.mcp.mcp_tools.df_query_mcp_tools import * -from flowcept.agents.mcp.mcp_tools.db_query_mcp_tools import * +Exports are resolved lazily (PEP 562). The MCP tool modules pull in optional +heavy dependencies (mcp, pandas, ...), and this package is also on the import +path of the harness capture hooks, which run as short-lived processes on an +interactive critical path and must stay stdlib-cheap. Importing +``flowcept.agents`` therefore imports nothing until an attribute is accessed. +""" + +_LAZY_MODULES = ( + "flowcept.agents.tool_result", + "flowcept.agents.mcp.mcp_tools", + "flowcept.agents.mcp.mcp_tools.df_query_mcp_tools", + "flowcept.agents.mcp.mcp_tools.db_query_mcp_tools", +) + + +def __getattr__(name): + import importlib + + for _mod_name in _LAZY_MODULES: + try: + _mod = importlib.import_module(_mod_name) + except Exception: + continue + if hasattr(_mod, name): + return getattr(_mod, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/src/flowcept/agents/academy/academy_plugin.py b/src/flowcept/agents/academy/academy_plugin.py new file mode 100644 index 00000000..761de0d9 --- /dev/null +++ b/src/flowcept/agents/academy/academy_plugin.py @@ -0,0 +1,1667 @@ +# academy_coscientist/plugins/flowcept_plugin.py +""" +Generic FlowCept provenance plugin for Academy-based applications. + +Designed around FlowCept's provenance model (TaskObject, WorkflowObject, +TelemetryCapture, Status). + +Provenance hierarchy produced: + WorkflowObject (one per run — top-level workflow) + └─ WorkflowObject (one per agent — sub-workflow, parent_workflow_id → top) + └─ TaskObject activity_id=action_name subtype=academy_action + └─ TaskObject activity_id=llm_call_type subtype=llm_call (parent_task_id) + └─ TaskObject activity_id=loop_name subtype=academy_loop + └─ TaskObject activity_id=agent_startup|shutdown subtype=academy_lifecycle + +Key FlowCept fields used correctly: + task_id — uuid per task record + workflow_id — per-agent sub-workflow id (linked to top-level via parent) + campaign_id — from Flowcept.campaign_id + parent_task_id — LLM tasks are children of the action that spawned them + group_id — loop events share a group_id for grouping + activity_id — action name, llm call type, loop name, lifecycle event + subtype — academy_action | academy_loop | academy_lifecycle | llm_call + used / generated — inputs and outputs + status — proper Status enum values (FINISHED | ERROR) + telemetry_at_start/end — CPU/memory snapshots via TelemetryCapture + node_name, hostname, etc. — via TaskObject.enrich_task_dict() + +LLM calls are linked to their parent action via a contextvars.ContextVar so that +any LLM call happening inside an async action is automatically a child of that action. + +Usage (zero agent code changes required): + + import my_llm_logging + plugin = FlowceptAcademyPlugin( + config={"workflow_name": "my-app"}, + llm_hook_register=my_llm_logging.register_llm_hook, + llm_hook_unregister=my_llm_logging.unregister_llm_hook, + ) + plugin.start() + try: + asyncio.run(my_academy_app()) + finally: + plugin.stop() + +LLM hook contract: + The hook function receives a dict payload with at minimum: + "type" — one of: parsed_json_result | chat_completion | + embed_result_local | embed_result_openai + "context" — dict of caller-supplied key/value tags (agent, call_type, ...) + "model" — model identifier string + All top-level fields from parsed JSON responses are automatically hoisted + into the generated record so they are directly queryable. +""" + +from __future__ import annotations + +import os + +import contextvars +import json as _json +import logging +import threading +import time +import uuid +from contextlib import contextmanager +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from concurrent.futures import ProcessPoolExecutor + + +# --------------------------------------------------------------------------- +# Provenance overhead timer +# --------------------------------------------------------------------------- + + +class _ProvenanceStats: + """ + Lightweight thread-safe accumulator for provenance capture timings. + + Records call counts, total/min/max wall-clock elapsed time (seconds) per + named *category*. All arithmetic uses ``time.perf_counter()`` so the + resolution is sub-microsecond on all supported platforms. + + Categories captured automatically: + ``action_emit`` — time spent inside ``_emit_action()`` + ``loop_emit`` — time spent inside ``_emit_loop_event()`` + ``lifecycle_emit``— time spent inside ``_emit_lifecycle()`` + ``llm_hook`` — time spent processing each LLM call in ``_on_llm_call()`` + ``intercept_task``— time spent inside ``AcademyInterceptor.intercept_task()`` + (serialisation + buffer append) + ``flush`` — wall-clock time for ``stop()`` + """ + + __slots__ = ("_lock", "_counts", "_totals", "_mins", "_maxs", "_raw") + + def __init__(self) -> None: + self._lock: threading.Lock = threading.Lock() + self._counts: dict[str, int] = {} + self._totals: dict[str, float] = {} + self._mins: dict[str, float] = {} + self._maxs: dict[str, float] = {} + # Raw per-event buffer: list of (timestamp_utc, category, elapsed_s) + self._raw: list[tuple[str, str, float]] = [] + + def record(self, category: str, elapsed: float) -> None: + """Add one observation for *category* (elapsed in seconds).""" + ts = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) + with self._lock: + if category not in self._counts: + self._counts[category] = 0 + self._totals[category] = 0.0 + self._mins[category] = float("inf") + self._maxs[category] = 0.0 + self._counts[category] += 1 + self._totals[category] += elapsed + if elapsed < self._mins[category]: + self._mins[category] = elapsed + if elapsed > self._maxs[category]: + self._maxs[category] = elapsed + self._raw.append((ts, category, elapsed)) + + def summary(self) -> str: + """Return a formatted table of all recorded categories.""" + col = 22 + header = f"{'Category':<{col}} {'N':>7} {'Total(ms)':>11} {'Mean(µs)':>9} {'Min(µs)':>8} {'Max(µs)':>8}" + sep = "-" * len(header) + rows = [header, sep] + with self._lock: + for cat in sorted(self._counts): + n = self._counts[cat] + total = self._totals[cat] + mean = (total / n) if n else 0.0 + mn = self._mins.get(cat, 0.0) + mx = self._maxs.get(cat, 0.0) + rows.append( + f"{cat:<{col}} {n:>7} {total * 1e3:>11.3f} {mean * 1e6:>9.1f} {mn * 1e6:>8.1f} {mx * 1e6:>8.1f}" + ) + return "\n".join(rows) + + def to_csv(self, path: str, workflow_id: str | None = None) -> None: + """Append one row per captured event to a CSV file (header written once). + + Each row is one individual provenance-capture observation so the file + can be used for distribution analysis, percentile queries, or time-series + plots without losing any detail. + + Columns + ------- + timestamp_utc, workflow_id, category, elapsed_us + """ + import csv + + write_header = not os.path.exists(path) + with self._lock: + raw_snapshot = list(self._raw) + + wf = workflow_id or "" + rows = [ + { + "timestamp_utc": ts, + "workflow_id": wf, + "category": cat, + "elapsed_us": round(elapsed * 1e6, 3), + } + for ts, cat, elapsed in raw_snapshot + ] + + fieldnames = ["timestamp_utc", "workflow_id", "category", "elapsed_us"] + with open(path, "a", newline="", encoding="utf-8") as fh: + writer = csv.DictWriter(fh, fieldnames=fieldnames) + if write_header: + writer.writeheader() + writer.writerows(rows) + + def reset(self) -> None: + with self._lock: + self._counts.clear() + self._totals.clear() + self._mins.clear() + self._maxs.clear() + self._raw.clear() + + +# Module-level singleton; activated when the plugin starts. +_PROV_STATS: _ProvenanceStats | None = None + + +@contextmanager +def _timed(category: str): + """Context manager: record wall-clock time for *category* in ``_PROV_STATS``.""" + t0 = time.perf_counter() + try: + yield + finally: + if _PROV_STATS is not None: + _PROV_STATS.record(category, time.perf_counter() - t0) + + +_log = logging.getLogger(__name__) + +# ContextVar: holds the task_id of the currently-executing Academy action (or +# loop) so that LLM calls made inside it can be registered as child tasks. +_current_action_task_id: contextvars.ContextVar[str | None] = contextvars.ContextVar( + "_current_action_task_id", default=None +) + +# ContextVar: holds the real Academy agent ID (e.g. "AgentId<82278eb6>") for +# the agent whose coroutine is currently executing. Set once per agent at +# startup so every LLM call emitted from within that agent carries the right ID. +_current_academy_agent_id: contextvars.ContextVar[str | None] = contextvars.ContextVar( + "_current_academy_agent_id", default=None +) + +# LLM payload types that carry a complete request+response pair. +_CAPTURE_LLM_TYPES = frozenset( + { + "parsed_json_result", + "chat_completion", + "embed_result_local", + "embed_result_openai", + } +) + + +# --------------------------------------------------------------------------- +# Academy interceptor — manages BaseInterceptor directly (Dask-style) +# --------------------------------------------------------------------------- + + +class AcademyInterceptor: + """ + Wrap FlowCept's BaseInterceptor directly (Dask-style). + + Exposes: + - start(workflow_name) — initialises the interceptor and MQDao + - stop() — flushes and closes the interceptor + - intercept_task(task_dict) — enriches with FlowCept standard fields and sends + - send_agent_workflow(...) — emits a WorkflowObject for an agent + - telemetry_capture — the underlying TelemetryCapture instance. + """ + + def __init__(self) -> None: + self._interceptor = None + self._workflow_id: str | None = None + self._campaign_id: str | None = None + + def start(self, workflow_name: str, campaign_id: str | None = None) -> None: + """Initialize the interceptor and emit the top-level workflow record.""" + from flowcept.flowceptor.adapters.base_interceptor import BaseInterceptor + from flowcept.commons.flowcept_dataclasses.workflow_object import WorkflowObject + + self._workflow_id = str(uuid.uuid4()) + self._campaign_id = campaign_id or str(uuid.uuid4()) + + # Create a fresh BaseInterceptor directly (Dask-style — no Flowcept controller) + self._interceptor = BaseInterceptor(kind="academy") + self._interceptor.start( + bundle_exec_id=self._workflow_id, + check_safe_stops=False, + ) + + # Emit the top-level workflow record + wf = WorkflowObject() + wf.workflow_id = self._workflow_id + wf.campaign_id = self._campaign_id + wf.name = workflow_name + self._interceptor.send_workflow_message(wf) + + def start_worker(self, workflow_id: str, campaign_id: str) -> None: + """ + Initialize this interceptor inside a worker process sharing an existing workflow. + + Unlike ``start()``, this method reuses the provided *workflow_id* and + *campaign_id* (already emitted by the parent process) so that all + provenance records from every worker process appear under the same + top-level WorkflowObject. No new WorkflowObject is emitted. + + Called automatically by ``_worker_init()``; use ``make_process_executor()`` + rather than calling this directly. + """ + from flowcept.flowceptor.adapters.base_interceptor import BaseInterceptor + + self._workflow_id = workflow_id + self._campaign_id = campaign_id + self._interceptor = BaseInterceptor(kind="academy") + self._interceptor.start( + bundle_exec_id=self._workflow_id, + check_safe_stops=False, + ) + + def stop(self) -> None: + """Flush and close the underlying interceptor.""" + if self._interceptor is None: + return + with _timed("flush"): + try: + self._interceptor.stop(check_safe_stops=False) + except Exception as e: + _log.warning("Interceptor stop error: %r", e) + self._interceptor = None + + @property + def telemetry_capture(self): + """Return the underlying TelemetryCapture instance, or None if not started.""" + return self._interceptor.telemetry_capture if self._interceptor else None + + def send_agent_workflow(self, agent_type: str, agent_id: str) -> str: + """Emit a WorkflowObject for an agent sub-workflow; return its workflow_id.""" + if self._interceptor is None: + return str(uuid.uuid4()) + from flowcept.commons.flowcept_dataclasses.workflow_object import WorkflowObject + + wf = WorkflowObject() + wf.workflow_id = str(uuid.uuid4()) + wf.name = f"{agent_type}:{agent_id}" + wf.campaign_id = self._campaign_id + wf.parent_workflow_id = self._workflow_id + wf.custom_metadata = {"agent_type": agent_type, "agent_id": agent_id} + self._interceptor.send_workflow_message(wf) + return wf.workflow_id + + def send_graph_workflow(self, graph_name: str, group_id: str) -> str: + """Emit a WorkflowObject for a LangGraph run sub-workflow; return its workflow_id. + + Same structure as send_agent_workflow but with graph-oriented metadata. + Allows AcademyInterceptor to serve as the shared interceptor for both + FlowceptAcademyPlugin and FlowceptLangGraphPlugin so that all provenance + records go into the same FlowCept buffer. + """ + if self._interceptor is None: + return str(uuid.uuid4()) + from flowcept.commons.flowcept_dataclasses.workflow_object import WorkflowObject + + wf = WorkflowObject() + wf.workflow_id = str(uuid.uuid4()) + wf.name = graph_name + wf.campaign_id = self._campaign_id + wf.parent_workflow_id = self._workflow_id + wf.custom_metadata = {"graph_name": graph_name, "group_id": group_id} + self._interceptor.send_workflow_message(wf) + return wf.workflow_id + + def intercept_task(self, task_dict: dict) -> None: + """Enrich task_dict with FlowCept standard fields and append to buffer.""" + if self._interceptor is None: + return + with _timed("intercept_task"): + from flowcept.commons.flowcept_dataclasses.task_object import TaskObject + from flowcept.commons.vocabulary import Status + + task_dict.setdefault("type", "task") + task_dict.setdefault("task_id", str(uuid.uuid4())) + task_dict.setdefault("workflow_id", self._workflow_id) + task_dict.setdefault("campaign_id", self._campaign_id) + + # Normalise status to proper enum value + raw = task_dict.get("status", "FINISHED") + if isinstance(raw, str): + try: + task_dict["status"] = Status[raw].value + except KeyError: + task_dict["status"] = Status.FINISHED.value + + # Enrich with node_name, login_name, hostname, public_ip, private_ip + TaskObject.enrich_task_dict(task_dict) + + self._interceptor.intercept(task_dict) + + +# --------------------------------------------------------------------------- +# Module-level singletons +# --------------------------------------------------------------------------- + +_PATCHER_INSTALLED: bool = False +_ACTIVE_INTERCEPTOR: AcademyInterceptor | None = None +_ORIG_PPE_INIT = None # holds the original ProcessPoolExecutor.__init__ while patched +_PERF_CSV_PATH: str | None = None # set by plugin.start(); forwarded to workers via initargs +_WORKER_PERF_CSV: str | None = None # set inside each worker by _worker_init + +# Per-agent sub-workflow IDs {agent_id_str -> sub_workflow_id} +_AGENT_WORKFLOWS: dict[str, str] = {} + +# Reverse map: Academy agent ID → agent class name, for enriching records that +# only have the academy ID (e.g. embed calls whose ctx has no "agent" key). +_AGENT_ID_TO_TYPE: dict[str, str] = {} + + +# --------------------------------------------------------------------------- +# Process-pool support — picklable worker initializer + executor factory +# --------------------------------------------------------------------------- + + +def _worker_init(workflow_id: str, campaign_id: str, perf_tracking: bool, perf_csv: str | None = None) -> None: + """ + ``ProcessPoolExecutor`` initializer for Flowcept provenance capture. + + Called automatically once per worker process before any agents run. + Sets up a fresh ``AcademyInterceptor`` that shares the parent's + *workflow_id* and *campaign_id* so all records appear in the same + provenance graph, patches ``academy.runtime.Runtime``, and registers + an atexit handler to flush the buffer when the worker exits. + + Do not call directly — use ``make_process_executor()`` instead. + """ + import atexit + + global _ACTIVE_INTERCEPTOR, _PROV_STATS, _WORKER_PERF_CSV + _WORKER_PERF_CSV = perf_csv + try: + interceptor = AcademyInterceptor() + interceptor.start_worker(workflow_id, campaign_id) + _ACTIVE_INTERCEPTOR = interceptor + _install_runtime_patches() + _PROV_STATS = _ProvenanceStats() if perf_tracking else None + atexit.register(_worker_shutdown) + except Exception as e: + _log.warning("[Flowcept] Worker process init failed: %r — provenance disabled in this worker.", e) + + +def _worker_shutdown() -> None: + """ + Atexit handler registered by ``_worker_init``. + + Flushes the worker's local provenance buffer to the MQ before the + process exits, ensuring no records are lost. Also appends this + worker's per-category timing rows to the shared perf CSV so that + action_emit / loop_emit / etc. from worker processes appear alongside + the main process's flush row. + """ + global _ACTIVE_INTERCEPTOR + interceptor = _ACTIVE_INTERCEPTOR + if interceptor is None: + return + wf_id = interceptor._workflow_id # save before stop() clears _interceptor + try: + interceptor.stop() + except Exception: + pass + _ACTIVE_INTERCEPTOR = None + if _PROV_STATS is not None and _WORKER_PERF_CSV: + try: + _PROV_STATS.to_csv(_WORKER_PERF_CSV, workflow_id=wf_id) + except Exception: + pass + + +def _patch_process_pool_executor() -> None: + """ + Monkey-patch ``ProcessPoolExecutor.__init__`` to inject ``_worker_init``. + + Any executor created while the Academy plugin is active automatically + receives ``_worker_init`` as its initializer. + + Only injects when no ``initializer`` is already provided by the caller, + so explicit user-supplied initializers are never overridden. + Called by ``FlowceptAcademyPlugin.start()``. + """ + global _ORIG_PPE_INIT + from concurrent.futures import ProcessPoolExecutor + + if _ORIG_PPE_INIT is not None: + return # already patched + + _ORIG_PPE_INIT = ProcessPoolExecutor.__init__ + _orig = _ORIG_PPE_INIT # capture for closure + + def _patched_init(self, max_workers=None, mp_context=None, initializer=None, initargs=(), **kw): + interceptor = _ACTIVE_INTERCEPTOR + if initializer is None and interceptor is not None: + initializer = _worker_init + initargs = ( + interceptor._workflow_id, + interceptor._campaign_id, + _PROV_STATS is not None, + _PERF_CSV_PATH, + ) + _orig(self, max_workers=max_workers, mp_context=mp_context, initializer=initializer, initargs=initargs, **kw) + + ProcessPoolExecutor.__init__ = _patched_init # type: ignore[method-assign] + + +def _unpatch_process_pool_executor() -> None: + """Restore the original ``ProcessPoolExecutor.__init__``. Called by ``stop()``.""" + global _ORIG_PPE_INIT + if _ORIG_PPE_INIT is None: + return + from concurrent.futures import ProcessPoolExecutor + + ProcessPoolExecutor.__init__ = _ORIG_PPE_INIT # type: ignore[method-assign] + _ORIG_PPE_INIT = None + + +def make_process_executor(max_workers: int | None = None) -> "ProcessPoolExecutor": + """ + Create a ``ProcessPoolExecutor`` with provenance capture pre-wired into every worker. + + Must be called while the Academy plugin is active — i.e., inside a + ``with Flowcept():`` block after the plugin has started — so that the + current ``workflow_id`` and ``campaign_id`` can be forwarded to workers. + + Each worker process receives its own ``AcademyInterceptor`` connected + to the same MQ backend. All task records carry the shared + *workflow_id* / *campaign_id*, so the full provenance graph is coherent + across processes. + + Parameters + ---------- + max_workers : int, optional + Number of worker processes. Defaults to ``os.cpu_count()``. + + Returns + ------- + ProcessPoolExecutor + Pass directly to ``Manager.from_exchange_factory(executors=executor)``. + + Raises + ------ + RuntimeError + If called before the Academy plugin has started. + + Example + ------- + :: + + from flowcept import Flowcept + from flowcept.agents.academy.academy_plugin import make_process_executor + + with Flowcept(): + executor = make_process_executor(max_workers=4) + async with await Manager.from_exchange_factory( + factory=exchange, executors=executor, + ) as manager: + counter = await manager.launch(CounterAgent) + ... + """ + from concurrent.futures import ProcessPoolExecutor + + interceptor = _ACTIVE_INTERCEPTOR + if interceptor is None: + raise RuntimeError( + "make_process_executor() must be called while the Flowcept Academy " + "plugin is running (inside a 'with Flowcept():' block)." + ) + return ProcessPoolExecutor( + max_workers=max_workers, + initializer=_worker_init, + initargs=( + interceptor._workflow_id, + interceptor._campaign_id, + _PROV_STATS is not None, + _PERF_CSV_PATH, + ), + ) + + +# --------------------------------------------------------------------------- +# Academy Runtime patcher +# --------------------------------------------------------------------------- + + +def _install_runtime_patches() -> None: + """ + Patch academy.runtime.Runtime at the class level. + + Every Runtime instance (one per agent) captures provenance without any + agent code changes. + """ + global _PATCHER_INSTALLED + if _PATCHER_INSTALLED: + return + + from academy.runtime import Runtime + + # ---- @action dispatch ------------------------------------------------ + _orig_action = Runtime.action + + async def _action_with_prov(self, action: str, source_id: Any, *, args: Any, kwargs: Any) -> Any: + interceptor = _ACTIVE_INTERCEPTOR + if interceptor is None: + return await _orig_action(self, action, source_id, args=args, kwargs=kwargs) + + tel_cap = interceptor.telemetry_capture + task_id = str(uuid.uuid4()) + tel_start = tel_cap.capture() if tel_cap else None + started_at = time.time() + + # Propagate action's task_id into context so nested LLM calls see it + token = _current_action_task_id.set(task_id) + error: BaseException | None = None + result: Any = None + try: + result = await _orig_action(self, action, source_id, args=args, kwargs=kwargs) + return result + except Exception as exc: + error = exc + raise + finally: + _current_action_task_id.reset(token) + ended_at = time.time() + tel_end = tel_cap.capture() if tel_cap else None + try: + _emit_action( + interceptor, + self, + action, + source_id, + args, + kwargs, + result, + error, + task_id, + started_at, + ended_at, + tel_start, + tel_end, + ) + except Exception: + pass + + Runtime.action = _action_with_prov # type: ignore[method-assign] + + # ---- @loop execution ------------------------------------------------- + _orig_execute_loop = Runtime._execute_loop + + async def _loop_with_prov(self, name: str, method: Any) -> None: + interceptor = _ACTIVE_INTERCEPTOR + if interceptor is None: + return await _orig_execute_loop(self, name, method) + + tel_cap = interceptor.telemetry_capture + task_id = str(uuid.uuid4()) + group_id = str(uuid.uuid4()) + tel_start = tel_cap.capture() if tel_cap else None + t0 = time.time() + _emit_loop_event(interceptor, self, name, "start", task_id, group_id, t0, tel_start) + + # Set context so LLM calls made inside the loop body are linked to + # this loop task. The action patch will override this for any + # @action dispatches that happen concurrently. + token = _current_action_task_id.set(task_id) + try: + await _orig_execute_loop(self, name, method) + tel_end = tel_cap.capture() if tel_cap else None + _emit_loop_event(interceptor, self, name, "exit", task_id, group_id, t0, tel_end) + except Exception as exc: + tel_end = tel_cap.capture() if tel_cap else None + _emit_loop_event(interceptor, self, name, "error", task_id, group_id, t0, tel_end, error=exc) + raise + finally: + _current_action_task_id.reset(token) + + Runtime._execute_loop = _loop_with_prov # type: ignore[method-assign] + + # ---- agent startup --------------------------------------------------- + _orig_start = Runtime._start + + async def _start_with_prov(self) -> None: + # CRITICAL: set the ContextVar BEFORE calling _orig_start. + # + # Academy's _start() creates all loop tasks and the exchange listener + # task via asyncio.create_task / spawn_guarded_background_task. Each + # new task gets a *copy* of the context at the moment of creation. + # If we set the ContextVar after _orig_start returns, those tasks + # already have stale context copies and will never see the agent ID. + # Setting it first ensures every task spawned during startup inherits + # the correct Academy agent ID automatically. + agent_type, agent_id = _agent_info(self) + _current_academy_agent_id.set(agent_id) + + await _orig_start(self) + + interceptor = _ACTIVE_INTERCEPTOR + if interceptor is None: + return + try: + sub_wf_id = interceptor.send_agent_workflow(agent_type, agent_id) + _AGENT_WORKFLOWS[agent_id] = sub_wf_id + _AGENT_ID_TO_TYPE[agent_id] = agent_type + _emit_lifecycle(interceptor, self, "agent_startup") + except Exception: + pass + + Runtime._start = _start_with_prov # type: ignore[method-assign] + + # ---- agent shutdown -------------------------------------------------- + _orig_shutdown = Runtime._shutdown + + async def _shutdown_with_prov(self) -> None: + interceptor = _ACTIVE_INTERCEPTOR + if interceptor is not None: + try: + _emit_lifecycle(interceptor, self, "agent_shutdown") + except Exception: + pass + await _orig_shutdown(self) + + Runtime._shutdown = _shutdown_with_prov # type: ignore[method-assign] + + _PATCHER_INSTALLED = True + + +def _uninstall_runtime_patches() -> None: + """Clear the active interceptor so all patches become no-ops.""" + global _ACTIVE_INTERCEPTOR + _ACTIVE_INTERCEPTOR = None + _AGENT_WORKFLOWS.clear() + _AGENT_ID_TO_TYPE.clear() + + +# --------------------------------------------------------------------------- +# Emit helpers +# --------------------------------------------------------------------------- + + +def _agent_info(runtime: Any) -> tuple[str, str]: + try: + agent_type = type(runtime.agent).__name__ + except Exception: + agent_type = "unknown" + try: + agent_id = str(runtime.agent_id) + except Exception: + agent_id = "unknown" + return agent_type, agent_id + + +def _agent_workflow_id(agent_id: str, fallback: str | None) -> str | None: + return _AGENT_WORKFLOWS.get(agent_id, fallback) + + +def _parse_source(source_id: Any) -> tuple[str | None, str | None, str | None]: + """ + Parse Academy source_id into (raw_str, agent_uid, source_workflow_id). + + source_id is either: + - UserId → a user/launcher, not an agent + - AgentId → another agent; look up its sub-workflow + Returns (source_str, source_agent_id, source_workflow_id). + """ + s = str(source_id) + if s.startswith("AgentId"): + wf_id = _AGENT_WORKFLOWS.get(s) + return s, s, wf_id + return s, None, None + + +def _tel_to_dict(tel: Any) -> dict | None: + if tel is None: + return None + try: + return tel.to_dict() + except Exception: + return None + + +def _emit_action( + interceptor: AcademyInterceptor, + runtime: Any, + action: str, + source_id: Any, + args: Any, + kwargs: Any, + result: Any, + error: BaseException | None, + task_id: str, + started_at: float, + ended_at: float, + tel_start: Any, + tel_end: Any, +) -> None: + with _timed("action_emit"): + agent_type, agent_id = _agent_info(runtime) + source_str, source_agent_id, source_workflow_id = _parse_source(source_id) + + custom: dict[str, Any] = { + "agent_type": agent_type, + "source_id": source_str, + # Always record whether this was a cross-agent call + "cross_agent_call": source_agent_id is not None, + } + if source_agent_id is not None: + # Explicit inter-agent provenance: who called this action and from + # which sub-workflow, so the full agent-to-agent graph is queryable. + custom["source_agent_id"] = source_agent_id + if source_workflow_id is not None: + custom["source_workflow_id"] = source_workflow_id + + task: dict[str, Any] = { + "task_id": task_id, + "subtype": "academy_action", + "activity_id": action, + "agent_id": agent_id, + "custom_metadata": custom, + "started_at": started_at, + "ended_at": ended_at, + "status": "ERROR" if error else "FINISHED", + "used": {"args": _safe_clip(args), "kwargs": _safe_clip(kwargs)}, + "generated": _safe_clip(result) if error is None else None, + "stderr": str(error) if error else None, + } + if tel_start is not None: + task["telemetry_at_start"] = _tel_to_dict(tel_start) + if tel_end is not None: + task["telemetry_at_end"] = _tel_to_dict(tel_end) + interceptor.intercept_task(task) + + +def _emit_loop_event( + interceptor: AcademyInterceptor, + runtime: Any, + loop_name: str, + event: str, + task_id: str, + group_id: str, + started_at: float, + tel: Any, + error: BaseException | None = None, +) -> None: + with _timed("loop_emit"): + agent_type, agent_id = _agent_info(runtime) + task: dict[str, Any] = { + "task_id": task_id, + "subtype": "academy_loop", + "activity_id": loop_name, + "group_id": group_id, + "agent_id": agent_id, + "custom_metadata": { + "agent_type": agent_type, + "loop_event": event, + }, + "started_at": started_at, + "ended_at": time.time(), + "status": "ERROR" if error else "FINISHED", + "stderr": str(error) if error else None, + } + tel_key = "telemetry_at_start" if event == "start" else "telemetry_at_end" + if tel is not None: + task[tel_key] = _tel_to_dict(tel) + interceptor.intercept_task(task) + + +def _emit_lifecycle(interceptor: AcademyInterceptor, runtime: Any, event: str) -> None: + with _timed("lifecycle_emit"): + agent_type, agent_id = _agent_info(runtime) + now = time.time() + task: dict[str, Any] = { + "subtype": "academy_lifecycle", + "activity_id": event, + "agent_id": agent_id, + "custom_metadata": {"agent_type": agent_type}, + "started_at": now, + "ended_at": now, + "status": "FINISHED", + } + interceptor.intercept_task(task) + + +# --------------------------------------------------------------------------- +# LLM hook — each LLM call becomes a child TaskObject of its parent action +# --------------------------------------------------------------------------- + + +def _on_llm_call(payload: dict) -> None: + interceptor = _ACTIVE_INTERCEPTOR + if interceptor is None: + return + call_type = payload.get("type", "") + if call_type not in _CAPTURE_LLM_TYPES: + return + with _timed("llm_hook"): + _process_llm_call(interceptor, payload, call_type) + + +def record_llm_call(payload: dict) -> None: + """ + Public API to record an LLM call into the active FlowCept provenance graph. + + Call this after any LLM invocation to capture it as a child TaskObject of + the enclosing Academy @action. The payload must include at minimum: + + type : "chat_completion" | "parsed_json_result" | + "embed_result_local" | "embed_result_openai" + model : str — model name requested + text : str — response text (for chat_completion) + usage : dict — {"prompt_tokens": int, "completion_tokens": int, "total_tokens": int} + + No-ops if the plugin has not been started. + """ + _on_llm_call(payload) + + +def openai_chat( + prompt: str, + model: str = "gpt-4o-mini", + system: str = "You are a helpful assistant.", + temperature: float | None = 0.3, + top_p: float | None = None, + max_tokens: int | None = None, + n: int = 1, + stop: list[str] | str | None = None, + frequency_penalty: float = 0.0, + presence_penalty: float = 0.0, + seed: int | None = None, + reasoning_effort: str | None = None, # "low" | "medium" | "high" for o1/o3 models + response_format: dict | None = None, + tools: list | None = None, + tool_choice: str | dict | None = None, + user: str | None = None, + context: dict | None = None, +) -> str: + """ + Make an OpenAI chat completion call and record it for FlowCept provenance. + + Captures all request parameters and response fields — including token usage + breakdown (reasoning tokens, cached tokens), system fingerprint, model + routing, finish reason, and tool calls — as a child TaskObject of the + enclosing Academy @action. + + Parameters + ---------- + prompt : str + The user message. + model : str + OpenAI model name (e.g. "gpt-4o-mini", "o1-mini", "o3"). + system : str + System prompt. + temperature : float or None + Sampling temperature. Set to None for reasoning models (o1/o3) that + do not accept a temperature parameter. + top_p : float or None + Nucleus sampling probability. + max_tokens : int or None + Maximum tokens to generate (maps to max_completion_tokens for o-series). + n : int + Number of completions to generate. + stop : list[str] or str or None + Stop sequences. + frequency_penalty : float + Penalises repeated tokens by frequency. + presence_penalty : float + Penalises tokens already present in the context. + seed : int or None + Seed for deterministic sampling. + reasoning_effort : str or None + "low" | "medium" | "high" — controls thinking budget on o1/o3 models. + response_format : dict or None + E.g. {"type": "json_object"} for JSON mode. + tools : list or None + OpenAI function-calling tool definitions. + tool_choice : str or dict or None + Tool selection strategy. + user : str or None + End-user identifier for abuse monitoring. + context : dict, optional + Extra tags stored in the provenance record (e.g. {"agent": "SummaryAgent"}). + + Returns + ------- + str + The model's response text (first choice). + """ + import openai as _openai + import time as _time + + client = _openai.OpenAI(api_key=os.environ.get("OPENAI_API_KEY")) + + messages = [ + {"role": "system", "content": system}, + {"role": "user", "content": prompt}, + ] + + # Build request kwargs — omit None values to avoid API errors on + # models that reject unsupported parameters (e.g. temperature for o1). + req: dict = {"model": model, "messages": messages, "n": n} + if temperature is not None: + req["temperature"] = temperature + if top_p is not None: + req["top_p"] = top_p + if max_tokens is not None: + req["max_completion_tokens"] = max_tokens + if stop is not None: + req["stop"] = stop + if frequency_penalty != 0.0: + req["frequency_penalty"] = frequency_penalty + if presence_penalty != 0.0: + req["presence_penalty"] = presence_penalty + if seed is not None: + req["seed"] = seed + if reasoning_effort is not None: + req["reasoning_effort"] = reasoning_effort + if response_format is not None: + req["response_format"] = response_format + if tools is not None: + req["tools"] = tools + if tool_choice is not None: + req["tool_choice"] = tool_choice + if user is not None: + req["user"] = user + + t0 = _time.time() + response = client.chat.completions.create(**req) + elapsed = _time.time() - t0 + + choice = response.choices[0] + text = choice.message.content or "" + usage = response.usage or {} + + # Detailed token breakdown (available on newer API versions) + usage_dict: dict = {} + if hasattr(usage, "prompt_tokens"): + usage_dict["prompt_tokens"] = usage.prompt_tokens + if hasattr(usage, "completion_tokens"): + usage_dict["completion_tokens"] = usage.completion_tokens + if hasattr(usage, "total_tokens"): + usage_dict["total_tokens"] = usage.total_tokens + # Reasoning tokens (o1/o3) + if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details: + ctd = usage.completion_tokens_details + usage_dict["reasoning_tokens"] = getattr(ctd, "reasoning_tokens", None) + usage_dict["accepted_prediction_tokens"] = getattr(ctd, "accepted_prediction_tokens", None) + usage_dict["rejected_prediction_tokens"] = getattr(ctd, "rejected_prediction_tokens", None) + # Cached prompt tokens + if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details: + ptd = usage.prompt_tokens_details + usage_dict["cached_tokens"] = getattr(ptd, "cached_tokens", None) + + # Tool calls in the response + tool_calls = None + if choice.message.tool_calls: + tool_calls = [ + { + "id": tc.id, + "type": tc.type, + "function": {"name": tc.function.name, "arguments": tc.function.arguments}, + } + for tc in choice.message.tool_calls + ] + + record_llm_call( + { + "type": "chat_completion", + # --- request --- + "model": model, + "model_used": response.model, + "messages": messages, + "user_prompt": prompt, + "system_prompt": system, + "temperature": temperature, + "temperature_suppressed": temperature is None, + "top_p": top_p, + "max_tokens": max_tokens, + "n": n, + "stop": stop, + "frequency_penalty": frequency_penalty, + "presence_penalty": presence_penalty, + "seed": seed, + "reasoning_effort": reasoning_effort, + "response_format": response_format, + "tools_provided": [t.get("function", {}).get("name") for t in (tools or [])], + "tool_choice": tool_choice, + # --- response --- + "text": text, + "finish_reason": choice.finish_reason, + "tool_calls": tool_calls, + "system_fingerprint": getattr(response, "system_fingerprint", None), + "response_id": response.id, + "created": response.created, + "usage": usage_dict, + "elapsed_s": elapsed, + # --- provenance tags --- + "context": context or {}, + } + ) + + return text + + +def anthropic_chat( + prompt: str, + model: str = "claude-haiku-4-5-20251001", + system: str = "You are a helpful assistant.", + max_tokens: int = 1024, + temperature: float | None = 1.0, + top_p: float | None = None, + top_k: int | None = None, + stop_sequences: list[str] | None = None, + tools: list | None = None, + tool_choice: dict | None = None, + thinking: dict | None = None, # {"type": "enabled", "budget_tokens": N} for extended thinking + metadata: dict | None = None, + context: dict | None = None, +) -> str: + """ + Make an Anthropic (Claude) chat completion call and record it for FlowCept provenance. + + Captures all request parameters and response fields — including input/output + token counts, cache usage, stop reason, thinking blocks, and tool use — as + a child TaskObject of the enclosing Academy @action. + + Parameters + ---------- + prompt : str + The user message. + model : str + Anthropic model ID (e.g. "claude-haiku-4-5-20251001", "claude-opus-4-6"). + system : str + System prompt. + max_tokens : int + Maximum tokens to generate (required by Anthropic API). + temperature : float or None + Sampling temperature (0–1). Set to None when using extended thinking. + top_p : float or None + Nucleus sampling probability. + top_k : int or None + Top-k sampling. + stop_sequences : list[str] or None + Custom stop sequences. + tools : list or None + Anthropic tool definitions. + tool_choice : dict or None + Tool selection strategy (e.g. {"type": "auto"}). + thinking : dict or None + Extended thinking config: {"type": "enabled", "budget_tokens": N}. + When set, temperature must be 1.0 or None. + metadata : dict or None + End-user metadata for abuse monitoring (e.g. {"user_id": "..."}). + context : dict, optional + Extra tags stored in the provenance record (e.g. {"agent": "SummaryAgent"}). + + Returns + ------- + str + The model's response text (first text block). + """ + import anthropic as _anthropic + import time as _time + + client = _anthropic.Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY")) + + req: dict = { + "model": model, + "max_tokens": max_tokens, + "system": system, + "messages": [{"role": "user", "content": prompt}], + } + if temperature is not None: + req["temperature"] = temperature + if top_p is not None: + req["top_p"] = top_p + if top_k is not None: + req["top_k"] = top_k + if stop_sequences: + req["stop_sequences"] = stop_sequences + if tools: + req["tools"] = tools + if tool_choice: + req["tool_choice"] = tool_choice + if thinking: + req["thinking"] = thinking + if metadata: + req["metadata"] = metadata + + t0 = _time.time() + response = client.messages.create(**req) + elapsed = _time.time() - t0 + + # Extract text and thinking blocks separately + text = "" + thinking_text = "" + tool_uses = [] + for block in response.content: + if block.type == "text": + text += block.text + elif block.type == "thinking": + thinking_text += getattr(block, "thinking", "") + elif block.type == "tool_use": + tool_uses.append( + { + "id": block.id, + "name": block.name, + "input": block.input, + } + ) + + usage = response.usage + usage_dict = { + "input_tokens": getattr(usage, "input_tokens", None), + "output_tokens": getattr(usage, "output_tokens", None), + "cache_creation_input_tokens": getattr(usage, "cache_creation_input_tokens", None), + "cache_read_input_tokens": getattr(usage, "cache_read_input_tokens", None), + } + + record_llm_call( + { + "type": "chat_completion", + # --- request --- + "model": model, + "model_used": response.model, + "messages": req["messages"], + "user_prompt": prompt, + "system_prompt": system, + "max_tokens": max_tokens, + "temperature": temperature, + "temperature_suppressed": temperature is None, + "top_p": top_p, + "top_k": top_k, + "stop_sequences": stop_sequences, + "thinking_budget_tokens": (thinking or {}).get("budget_tokens"), + "tools_provided": [t.get("name") for t in (tools or [])], + "tool_choice": tool_choice, + # --- response --- + "text": text, + "thinking_text": thinking_text if thinking_text else None, + "finish_reason": response.stop_reason, + "stop_sequence": response.stop_sequence, + "tool_uses": tool_uses if tool_uses else None, + "response_id": response.id, + "usage": usage_dict, + "elapsed_s": elapsed, + # --- provenance tags --- + "context": context or {}, + } + ) + + return text + + +# --------------------------------------------------------------------------- +# FlowceptAnthropicClient — wraps anthropic.Anthropic / AsyncAnthropic to +# capture full provenance for every messages.create / messages.stream call. +# --------------------------------------------------------------------------- + + +class FlowceptAnthropicClient: + """ + Wrap an ``anthropic.Anthropic`` (or ``AsyncAnthropic``) client for provenance capture. + + Records every ``messages.create`` / ``messages.stream`` call as a FlowCept + provenance record (subtype=llm_call) via ``record_llm_call()``. + + Usage:: + + import anthropic + from flowcept.agents.academy.academy_plugin import FlowceptAnthropicClient + + client = FlowceptAnthropicClient(anthropic.Anthropic(), agent_name="my-agent") + response = client.messages.create( + model="claude-haiku-4-5-20251001", + max_tokens=1024, + messages=[{"role": "user", "content": "Hello"}], + ) + """ + + def __init__(self, client, agent_name=None, context=None): + self._inner = client + self._agent_name = agent_name + self._context: dict = context or {} + self.messages = _FlowceptAnthropicMessages(client.messages, agent_name, self._context) + + def __getattr__(self, name): + """Delegate attribute access to the wrapped client.""" + return getattr(self._inner, name) + + +class _FlowceptAnthropicMessages: + def __init__(self, messages_resource, agent_name, context): + self._inner = messages_resource + self._agent_name = agent_name + self._context = context + + def __getattr__(self, name): + return getattr(self._inner, name) + + def create(self, **kwargs): + import time as _time + + t0 = _time.time() + result = self._inner.create(**kwargs) + self._record(kwargs, result, _time.time() - t0) + return result + + async def async_create(self, **kwargs): + import time as _time + + t0 = _time.time() + result = await self._inner.create(**kwargs) + self._record(kwargs, result, _time.time() - t0) + return result + + def stream(self, **kwargs): + import time as _time + + return _FlowceptAnthropicStream(self._inner.stream(**kwargs), kwargs, self._record, _time.time()) + + def _record(self, kwargs, result, elapsed): + try: + model = kwargs.get("model", "unknown") + text = "" + thinking_text = "" + tool_uses = [] + for block in getattr(result, "content", []): + btype = getattr(block, "type", None) + if btype == "text": + text += getattr(block, "text", "") + elif btype == "thinking": + thinking_text += getattr(block, "thinking", "") + elif btype == "tool_use": + tool_uses.append( + { + "id": getattr(block, "id", None), + "name": getattr(block, "name", None), + "input": getattr(block, "input", None), + } + ) + usage = getattr(result, "usage", None) + usage_dict = ( + { + attr: getattr(usage, attr, None) + for attr in ( + "input_tokens", + "output_tokens", + "cache_creation_input_tokens", + "cache_read_input_tokens", + ) + } + if usage + else {} + ) + ctx = dict(self._context) + if self._agent_name: + ctx["agent_name"] = self._agent_name + payload = { + "type": "chat_completion", + "model": model, + "model_used": getattr(result, "model", model), + "messages": kwargs.get("messages"), + "system_prompt": kwargs.get("system"), + "max_tokens": kwargs.get("max_tokens"), + "temperature": kwargs.get("temperature"), + "top_p": kwargs.get("top_p"), + "top_k": kwargs.get("top_k"), + "stop_sequences": kwargs.get("stop_sequences"), + "thinking_budget_tokens": (kwargs.get("thinking") or {}).get("budget_tokens"), + "tools_provided": [t.get("name") for t in (kwargs.get("tools") or [])], + "tool_choice": kwargs.get("tool_choice"), + "text": text, + "thinking_text": thinking_text or None, + "finish_reason": getattr(result, "stop_reason", None), + "stop_sequence": getattr(result, "stop_sequence", None), + "tool_uses": tool_uses or None, + "response_id": getattr(result, "id", None), + "usage": usage_dict, + "elapsed_s": elapsed, + "context": ctx, + } + record_llm_call({k: v for k, v in payload.items() if v is not None}) + except Exception: + pass + + +class _FlowceptAnthropicStream: + def __init__(self, ctx_mgr, kwargs, record_fn, t0): + self._ctx_mgr = ctx_mgr + self._kwargs = kwargs + self._record_fn = record_fn + self._t0 = t0 + self._stream = None + + def __enter__(self): + self._stream = self._ctx_mgr.__enter__() + return self._stream + + def __exit__(self, *args): + import time as _time + + result = None + try: + result = self._stream.get_final_message() + except Exception: + pass + if result is not None: + self._record_fn(self._kwargs, result, _time.time() - self._t0) + return self._ctx_mgr.__exit__(*args) + + +def _process_llm_call(interceptor: AcademyInterceptor, payload: dict, call_type: str) -> None: + """Inner body of ``_on_llm_call``, extracted so ``_timed`` wraps everything.""" + now = time.time() + ctx = payload.get("context") or {} + model = payload.get("model", "unknown") + model_used = payload.get("model_used", model) + has_error = "error" in payload + + # parent_task_id: links this LLM call to the enclosing action or loop task + parent_task_id = _current_action_task_id.get(None) + + # agent_id: always use the real Academy ID from the ContextVar (set before + # task creation so all tasks inherit it). ctx["agent"] is the class name + # logged by application code — useful but not a reliable unique identifier. + academy_agent_id = _current_academy_agent_id.get(None) + ctx_agent_name = ctx.get("agent") or ( + _AGENT_ID_TO_TYPE.get(academy_agent_id, "unknown") if academy_agent_id else "unknown" + ) + # Use the real Academy ID as the primary agent_id in provenance records. + # Fall back to the class name only when the ID is genuinely unavailable. + agent_id = academy_agent_id or ctx_agent_name + + # --- temperature --- + temp_requested = payload.get("temperature_requested", payload.get("temperature")) + temp_sent = payload.get("temperature_sent", payload.get("temperature")) + temp_suppressed = payload.get("temperature_suppressed", False) + + if temp_requested is None: + if temp_suppressed: + temp_null_reason = "reasoning_model_does_not_accept_temperature" + elif call_type in ("embed_result_local", "embed_result_openai"): + temp_null_reason = "not_applicable_for_embeddings" + else: + temp_null_reason = "not_set" + else: + temp_null_reason = None + + # --- used (request / inputs) --- + used: dict[str, Any] = { + "model": model, + "model_used": model_used, + "call_type": ctx.get("call_type") or call_type, + "temperature_requested": temp_requested, + "temperature_sent": temp_sent, + "temperature_suppressed": temp_suppressed, + # Always store both the class-name agent and the Academy agent ID so + # the record is self-contained for provenance queries. + "agent_class": ctx_agent_name, + "academy_agent_id": academy_agent_id, + } + if temp_null_reason is not None: + used["temperature_null_reason"] = temp_null_reason + if "messages" in payload: + used["messages"] = payload["messages"] + if "system_instructions" in payload: + used["system_instructions"] = payload["system_instructions"] + if "user_prompt" in payload: + used["user_prompt"] = payload["user_prompt"] + if "schema_hint" in payload: + used["schema_hint"] = payload["schema_hint"] + for k, v in ctx.items(): + used.setdefault(f"ctx_{k}", v) + + # --- generated (response / outputs) --- + usage = payload.get("usage") or {} + generated: dict[str, Any] = { + "finish_reason": payload.get("finish_reason"), + "prompt_tokens": usage.get("prompt_tokens"), + "completion_tokens": usage.get("completion_tokens"), + "total_tokens": usage.get("total_tokens"), + } + + if has_error: + generated["error"] = str(payload["error"]) + + elif call_type == "parsed_json_result": + # _call_llm_json path: parsed_response contains the full structured output. + parsed_resp = payload.get("parsed_response", {}) + generated["parsed_response"] = parsed_resp + generated["fallback_used"] = payload.get("fallback_used", False) + generated["raw_response_text"] = payload.get("raw_response_text", "") + # Hoist all top-level fields from the parsed response so they are + # directly queryable without unwrapping parsed_response. + if isinstance(parsed_resp, dict): + for k, v in parsed_resp.items(): + generated.setdefault(k, v) + + elif call_type == "chat_completion": + response_text = payload.get("text", "") + generated["response_text"] = response_text + + # Try to parse JSON embedded in the response text. + parsed_inline: dict | None = None + if response_text: + s = response_text.strip() + if s.startswith("```"): + lines = s.splitlines() + s = "\n".join(lines[1:-1]).strip() + try: + parsed_inline = _json.loads(s) + except Exception: + start = min( + (s.find("{") if s.find("{") != -1 else len(s)), + (s.find("[") if s.find("[") != -1 else len(s)), + ) + end = max(s.rfind("}"), s.rfind("]")) + if 0 <= start < end: + try: + parsed_inline = _json.loads(s[start : end + 1]) + except Exception: + pass + + if isinstance(parsed_inline, dict): + generated["parsed_response"] = parsed_inline + # Hoist all top-level fields so they are directly queryable. + for k, v in parsed_inline.items(): + generated.setdefault(k, v) + + elif call_type in ("embed_result_local", "embed_result_openai"): + generated["num_vectors"] = payload.get("num_vectors") + generated["embed_dim"] = payload.get("dim") + + task: dict[str, Any] = { + "subtype": "llm_call", + "activity_id": ctx.get("call_type") or call_type, + # agent_id uses the real Academy ID when available + "agent_id": agent_id, + "custom_metadata": { + "llm_call_type": call_type, + "hyp_id": ctx.get("hyp_id"), + }, + "started_at": now, + "ended_at": now, + "status": "ERROR" if has_error else "FINISHED", + "used": used, + "generated": generated, + } + if parent_task_id is not None: + task["parent_task_id"] = parent_task_id + + interceptor.intercept_task(task) + + +# --------------------------------------------------------------------------- +# Public plugin class +# --------------------------------------------------------------------------- + + +class FlowceptAcademyPlugin: + """ + Generic FlowCept provenance plugin for any Academy-based application. + + Patches Academy's Runtime at the class level (zero agent code changes required). + Produces a full provenance graph: campaign → workflow → agent sub-workflows → + action tasks → LLM child tasks, with CPU/memory telemetry and node enrichment. + + Parameters + ---------- + config : dict, optional + Plugin configuration keys: + enabled (bool, default True) + workflow_name (str, default "academy-workflow") + llm_hook_register : callable, optional + Called as llm_hook_register(_on_llm_call) on start to wire up LLM telemetry. + Should match the signature used by your application's logging/hook module. + llm_hook_unregister : callable, optional + Called as llm_hook_unregister(_on_llm_call) on stop to remove the hook. + """ + + def __init__( + self, + config: dict | None = None, + llm_hook_register: Any = None, + llm_hook_unregister: Any = None, + ) -> None: + cfg = config or {} + self._enabled: bool = cfg.get("enabled", False) + self._workflow_name: str = cfg.get("workflow_name", "academy-workflow") + self._campaign_id: str | None = cfg.get("campaign_id", None) + self._perf_tracking: bool = cfg.get("performance_tracking", True) + self._perf_csv: str | None = cfg.get("perf_csv", None) # explicit override + self._llm_hook_register = llm_hook_register + self._llm_hook_unregister = llm_hook_unregister + self._interceptor = AcademyInterceptor() + self._started = False + + def start(self) -> "FlowceptAcademyPlugin": + """Start provenance capture: install patches, hooks, and the interceptor.""" + if not self._enabled or self._started: + return self + global _ACTIVE_INTERCEPTOR, _PROV_STATS, _PERF_CSV_PATH + try: + _PROV_STATS = _ProvenanceStats() if self._perf_tracking else None + self._interceptor.start(self._workflow_name, campaign_id=self._campaign_id) + self._campaign_id = self._interceptor._campaign_id + _ACTIVE_INTERCEPTOR = self._interceptor + wf_id = self._interceptor._workflow_id + _PERF_CSV_PATH = self._perf_csv or f"provenance_perf_{wf_id}.csv" + _install_runtime_patches() + _patch_process_pool_executor() + if self._llm_hook_register is not None: + self._llm_hook_register(_on_llm_call) + self._started = True + campaign_id = self._interceptor._campaign_id + llm_status = "enabled" if self._llm_hook_register else "disabled (no hook provided)" + print( + f"[FlowceptAcademyPlugin] Started\n" + f" workflow_id : {wf_id}\n" + f" campaign_id : {campaign_id}\n" + f" Capturing : agent lifecycle (sub-workflows), action dispatch " + f"(telemetry + parent_task_id), loop events (group_id).\n" + f" LLM capture : {llm_status}", + flush=True, + ) + except Exception as e: + print( + f"[FlowceptAcademyPlugin] WARNING: failed to start — {e!r}. Continuing without provenance capture.", + flush=True, + ) + _log.exception("FlowceptAcademyPlugin start failed") + self._enabled = False + return self + + def stop(self) -> None: + """Stop provenance capture: unregister hooks, remove patches, flush records.""" + if not self._started: + return + if self._llm_hook_unregister is not None: + self._llm_hook_unregister(_on_llm_call) + _uninstall_runtime_patches() + _unpatch_process_pool_executor() + try: + self._interceptor.stop() + except Exception as e: + print(f"[FlowceptAcademyPlugin] Warning during stop: {e!r}", flush=True) + self._started = False + print("[FlowceptAcademyPlugin] Stopped.", flush=True) + if _PROV_STATS is not None: + print( + "\n[FlowceptAcademyPlugin] Provenance overhead report:\n" + + _PROV_STATS.summary() + + "\n (N = event count; Total/Mean/Min/Max in ms/µs respectively)\n", + flush=True, + ) + # Derive CSV path: explicit config > default based on workflow_id + wf_id = self._interceptor._workflow_id + csv_path = self._perf_csv or f"provenance_perf_{wf_id}.csv" + try: + _PROV_STATS.to_csv(csv_path, workflow_id=wf_id) + print( + f"[FlowceptAcademyPlugin] Performance stats written to {csv_path}", + flush=True, + ) + except Exception as e: + print( + f"[FlowceptAcademyPlugin] Warning: could not write perf CSV — {e!r}", + flush=True, + ) + + +# Keep old name working +FlowceptPlugin = FlowceptAcademyPlugin + + +# --------------------------------------------------------------------------- +# Serialisation helper +# --------------------------------------------------------------------------- + + +def _safe_clip(obj: Any, _depth: int = 0) -> Any: + """Recursively convert objects to JSON-serialisable form without truncation.""" + if _depth > 8: + return str(obj) + if obj is None or isinstance(obj, (bool, int, float, str)): + return obj + if isinstance(obj, dict): + return {str(k): _safe_clip(v, _depth + 1) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return [_safe_clip(v, _depth + 1) for v in obj] + return repr(obj) diff --git a/src/flowcept/agents/autogen/autogen_plugin.py b/src/flowcept/agents/autogen/autogen_plugin.py new file mode 100644 index 00000000..837e7d86 --- /dev/null +++ b/src/flowcept/agents/autogen/autogen_plugin.py @@ -0,0 +1,1392 @@ +# academy_coscientist/plugins/flowcept_autogen_plugin.py +""" +FlowCept provenance plugin for AutoGen (autogen_agentchat 0.7+) workflows. + +Captures provenance for AutoGen team runs by consuming the ``run_stream()`` +generator and recording every message as a FlowCept TaskObject, with the +overall team run as a WorkflowObject. + +Provenance hierarchy produced: + WorkflowObject (one per team.run() call) + └─ TaskObject subtype=autogen_run activity_id= + └─ TaskObject subtype=autogen_message activity_id= + (one record per message yielded by run_stream) + +Key FlowCept fields: + task_id — uuid per event + workflow_id — global shared workflow id (setdefault pattern) + campaign_id — from Flowcept.campaign_id + parent_task_id — messages are children of the enclosing run task + group_id — all tasks within one team.run() share a group_id + activity_id — team name | agent name | model name + subtype — autogen_run | autogen_message | autogen_result + used / generated — message content, source, recipient + status — FINISHED | ERROR + +Usage +----- +Standalone:: + + plugin = FlowceptAutoGenPlugin(config={"workflow_name": "my-team"}) + plugin.start() + result = asyncio.run(plugin.run_team(team, "Do something useful")) + plugin.stop() + +Or as a context manager:: + + with FlowceptAutoGenPlugin(config={"workflow_name": "my-team"}) as plugin: + result = asyncio.run(plugin.run_team(team, "task")) + +Shared with Academy plugin:: + + academy_plugin = FlowceptAcademyPlugin(config={...}).start() + autogen_plugin = FlowceptAutoGenPlugin.from_academy_plugin(academy_plugin) + result = asyncio.run(autogen_plugin.run_team(team, "task")) + academy_plugin.stop() # flushes the shared buffer +""" + +from __future__ import annotations + +import os + +import contextvars +import time +import uuid +import threading +import logging +from contextlib import contextmanager +from typing import Any + +_log = logging.getLogger(__name__) + +# ContextVar: holds the task_id of the currently-executing AutoGen message (or +# run). record_llm_call() reads this to set parent_task_id on LLM call records. +_current_autogen_task_id: contextvars.ContextVar[str | None] = contextvars.ContextVar( + "_current_autogen_task_id", default=None +) + +# ContextVar: holds the agent name (source) currently being processed, so that +# record_llm_call() can tag LLM records with the agent that made the call. +_current_autogen_agent_id: contextvars.ContextVar[str | None] = contextvars.ContextVar( + "_current_autogen_agent_id", default=None +) + + +# --------------------------------------------------------------------------- +# Provenance overhead timer +# --------------------------------------------------------------------------- + + +class _ProvenanceStats: + __slots__ = ("_lock", "_counts", "_totals", "_mins", "_maxs", "_raw") + + def __init__(self) -> None: + self._lock: threading.Lock = threading.Lock() + self._counts: dict[str, int] = {} + self._totals: dict[str, float] = {} + self._mins: dict[str, float] = {} + self._maxs: dict[str, float] = {} + self._raw: list[tuple[str, str, float]] = [] + + def record(self, category: str, elapsed: float) -> None: + ts = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) + with self._lock: + if category not in self._counts: + self._counts[category] = 0 + self._totals[category] = 0.0 + self._mins[category] = float("inf") + self._maxs[category] = 0.0 + self._counts[category] += 1 + self._totals[category] += elapsed + if elapsed < self._mins[category]: + self._mins[category] = elapsed + if elapsed > self._maxs[category]: + self._maxs[category] = elapsed + self._raw.append((ts, category, elapsed)) + + def summary(self) -> str: + col = 22 + header = f"{'Category':<{col}} {'N':>7} {'Total(ms)':>11} {'Mean(µs)':>9} {'Min(µs)':>8} {'Max(µs)':>8}" + sep = "-" * len(header) + rows = [header, sep] + with self._lock: + for cat in sorted(self._counts): + n = self._counts[cat] + total = self._totals[cat] + mean = (total / n) if n else 0.0 + mn = self._mins.get(cat, 0.0) + mx = self._maxs.get(cat, 0.0) + rows.append( + f"{cat:<{col}} {n:>7} {total * 1e3:>11.3f} {mean * 1e6:>9.1f} {mn * 1e6:>8.1f} {mx * 1e6:>8.1f}" + ) + return "\n".join(rows) + + def to_csv(self, path: str, workflow_id: str | None = None) -> None: + import csv + + write_header = not os.path.exists(path) + with self._lock: + raw_snapshot = list(self._raw) + wf = workflow_id or "" + rows = [ + { + "timestamp_utc": ts, + "workflow_id": wf, + "category": cat, + "elapsed_us": round(elapsed * 1e6, 3), + } + for ts, cat, elapsed in raw_snapshot + ] + fieldnames = ["timestamp_utc", "workflow_id", "category", "elapsed_us"] + with open(path, "a", newline="", encoding="utf-8") as fh: + writer = csv.DictWriter(fh, fieldnames=fieldnames) + if write_header: + writer.writeheader() + writer.writerows(rows) + + +# --------------------------------------------------------------------------- +# Standalone interceptor wrapper +# --------------------------------------------------------------------------- + + +class _AutoGenInterceptor: + """Standalone FlowCept interceptor for AutoGen provenance (Dask-style).""" + + def __init__(self) -> None: + self._interceptor = None + self._workflow_id: str | None = None + self._campaign_id: str | None = None + + def start(self, workflow_name: str, campaign_id: str | None = None) -> None: + from flowcept.flowceptor.adapters.base_interceptor import BaseInterceptor + from flowcept.commons.flowcept_dataclasses.workflow_object import WorkflowObject + + self._workflow_id = str(uuid.uuid4()) + self._campaign_id = campaign_id or str(uuid.uuid4()) + + self._interceptor = BaseInterceptor(kind="autogen") + self._interceptor.start( + bundle_exec_id=self._workflow_id, + check_safe_stops=False, + ) + + wf = WorkflowObject() + wf.workflow_id = self._workflow_id + wf.campaign_id = self._campaign_id + wf.name = workflow_name + self._interceptor.send_workflow_message(wf) + + def stop(self) -> None: + if self._interceptor is None: + return + try: + self._interceptor.stop(check_safe_stops=False) + except Exception as e: + _log.warning("Interceptor stop error: %r", e) + self._interceptor = None + + def send_team_workflow(self, team_name: str, group_id: str) -> str: + if self._interceptor is None: + return str(uuid.uuid4()) + from flowcept.commons.flowcept_dataclasses.workflow_object import WorkflowObject + + wf = WorkflowObject() + wf.workflow_id = str(uuid.uuid4()) + wf.name = team_name + wf.campaign_id = self._campaign_id + wf.parent_workflow_id = self._workflow_id + wf.custom_metadata = {"group_id": group_id, "framework": "autogen"} + self._interceptor.send_workflow_message(wf) + return wf.workflow_id + + # Accept the same name used by AcademyInterceptor so shared interceptors work + send_graph_workflow = send_team_workflow + + def intercept_task(self, task_dict: dict) -> None: + if self._interceptor is None: + return + from flowcept.commons.flowcept_dataclasses.task_object import TaskObject + from flowcept.commons.vocabulary import Status + + task_dict.setdefault("task_id", str(uuid.uuid4())) + task_dict.setdefault("workflow_id", self._workflow_id) + task_dict.setdefault("campaign_id", self._campaign_id) + + raw = task_dict.get("status", "FINISHED") + if isinstance(raw, str): + try: + task_dict["status"] = Status[raw].value + except KeyError: + task_dict["status"] = Status.FINISHED.value + + TaskObject.enrich_task_dict(task_dict) + self._interceptor.intercept(task_dict) + + +# --------------------------------------------------------------------------- +# Provenance stream runner +# --------------------------------------------------------------------------- + + +def _safe_clip(obj: Any, depth: int = 0) -> Any: + if depth > 6: + return str(obj) + if obj is None or isinstance(obj, (bool, int, float, str)): + return obj + if isinstance(obj, dict): + return {str(k): _safe_clip(v, depth + 1) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return [_safe_clip(v, depth + 1) for v in obj] + # AutoGen message objects + try: + if hasattr(obj, "model_dump"): + return _safe_clip(obj.model_dump(), depth + 1) + except Exception: + pass + try: + if hasattr(obj, "__dict__"): + return _safe_clip(vars(obj), depth + 1) + except Exception: + pass + return repr(obj) + + +# --------------------------------------------------------------------------- +# Automatic LLM patching — mirrors academy's _install_runtime_patches() +# --------------------------------------------------------------------------- + +_orig_assistant_init = None +_patches_installed = False + + +def _install_autogen_patches() -> None: + """Patch AssistantAgent.__init__ to auto-wrap _model_client with FlowceptModelClient.""" + global _orig_assistant_init, _patches_installed + if _patches_installed: + return + try: + from autogen_agentchat.agents import AssistantAgent + + _orig_assistant_init = AssistantAgent.__init__ + + def _patched_init(self_agent, *args, **kwargs): + _orig_assistant_init(self_agent, *args, **kwargs) + mc = getattr(self_agent, "_model_client", None) + if mc is not None and not isinstance(mc, FlowceptModelClient): + self_agent._model_client = FlowceptModelClient(mc, agent_name=getattr(self_agent, "name", None)) + + AssistantAgent.__init__ = _patched_init + _patches_installed = True + except (ImportError, AttributeError): + pass + + +def _uninstall_autogen_patches() -> None: + """Restore original AssistantAgent.__init__.""" + global _orig_assistant_init, _patches_installed + if not _patches_installed: + return + try: + from autogen_agentchat.agents import AssistantAgent + + AssistantAgent.__init__ = _orig_assistant_init + _orig_assistant_init = None + _patches_installed = False + except (ImportError, AttributeError): + pass + + +def _ensure_agents_wrapped(team: Any) -> None: + """Walk all participants of a team and wrap their _model_client if not already wrapped. + + Handles agents created before the plugin was started. + """ + participants = getattr(team, "_participants", None) or getattr(team, "agents", None) or [] + for agent in participants: + mc = getattr(agent, "_model_client", None) + if mc is not None and not isinstance(mc, FlowceptModelClient): + agent._model_client = FlowceptModelClient(mc, agent_name=getattr(agent, "name", None)) + + +async def _run_with_provenance( + team: Any, + task: str, + interceptor: Any, + stats: _ProvenanceStats | None, + team_name: str = "autogen_team", + source_agent_id: str | None = None, +) -> Any: + """ + Consume team.run_stream() and emit provenance records. + + Emits a FlowCept provenance record for each message plus one overall run + record. Returns the final TaskResult. + """ + from autogen_agentchat.base import TaskResult + from autogen_core import CancellationToken + + # Wrap any pre-existing agents' model clients (created before plugin start) + _ensure_agents_wrapped(team) + + group_id = str(uuid.uuid4()) + run_task_id = str(uuid.uuid4()) + run_start = time.time() + + # Emit a sub-WorkflowObject for this team run + with _timed("send_graph_workflow"): + interceptor.send_graph_workflow(team_name, group_id) + + custom_meta: dict = {"team_name": team_name, "framework": "autogen"} + if source_agent_id: + custom_meta["source_agent_id"] = source_agent_id + + messages_captured: list[dict] = [] + final_result: Any = None + + # Set contextvar to run_task_id so record_llm_call() links LLM calls to this run + _token_task = _current_autogen_task_id.set(run_task_id) + _token_agent = _current_autogen_agent_id.set(team_name) + + t_stream_start = time.perf_counter() + + try: + stream = team.run_stream(task=task, cancellation_token=CancellationToken()) + async for item in stream: + t0 = time.perf_counter() + if isinstance(item, TaskResult): + final_result = item + else: + # Each item is a BaseChatMessage or BaseAgentEvent + msg_task_id = str(uuid.uuid4()) + source = getattr(item, "source", None) or "unknown" + content = getattr(item, "content", None) + msg_type = type(item).__name__ + + # Update contextvar so LLM calls within this message are children of it + _current_autogen_task_id.set(msg_task_id) + _current_autogen_agent_id.set(source) + + # Serialize content + if isinstance(content, str): + content_clip = content + else: + content_clip = _safe_clip(content) + + msg_record: dict = { + "task_id": msg_task_id, + "subtype": "autogen_message", + "activity_id": source, + "group_id": group_id, + "parent_task_id": run_task_id, + "started_at": time.time(), + "ended_at": time.time(), + "status": "FINISHED", + "used": {"task": task, "agent": source}, + "generated": {"content": content_clip, "message_type": msg_type}, + "custom_metadata": { + "agent_name": source, + "message_type": msg_type, + "framework": "autogen", + }, + } + interceptor.intercept_task(msg_record) + messages_captured.append( + { + "source": source, + "content": content_clip[:200] if isinstance(content_clip, str) else content_clip, + } + ) + if stats is not None: + stats.record("message_intercept", time.perf_counter() - t0) + + except Exception as exc: + # Emit the run record as ERROR then re-raise + run_task: dict = { + "task_id": run_task_id, + "subtype": "autogen_run", + "activity_id": team_name, + "group_id": group_id, + "started_at": run_start, + "ended_at": time.time(), + "status": "ERROR", + "used": {"task": task}, + "generated": {"messages": messages_captured}, + "stderr": str(exc), + "custom_metadata": custom_meta, + } + with _timed("run_emit"): + interceptor.intercept_task(run_task) + _current_autogen_task_id.reset(_token_task) + _current_autogen_agent_id.reset(_token_agent) + raise + + _current_autogen_task_id.reset(_token_task) + _current_autogen_agent_id.reset(_token_agent) + + if stats is not None: + stats.record("stream_total", time.perf_counter() - t_stream_start) + + # Emit the overall run record (ONE complete record with inputs + outputs) + summary = _safe_clip(getattr(final_result, "stop_reason", None) or "completed") + msg_count = len(messages_captured) + last_msg = messages_captured[-1]["content"] if messages_captured else "" + + run_task = { + "task_id": run_task_id, + "subtype": "autogen_run", + "activity_id": team_name, + "group_id": group_id, + "started_at": run_start, + "ended_at": time.time(), + "status": "FINISHED", + "used": {"task": task}, + "generated": { + "stop_reason": summary, + "message_count": msg_count, + "last_message": last_msg, + "messages": messages_captured, + }, + "custom_metadata": custom_meta, + } + with _timed("run_emit"): + interceptor.intercept_task(run_task) + + return final_result + + +# --------------------------------------------------------------------------- +# Module-level active interceptor — set by start() / from_academy_plugin() +# --------------------------------------------------------------------------- + +_ACTIVE_INTERCEPTOR = None +_PROV_STATS: _ProvenanceStats | None = None + + +@contextmanager +def _timed(category: str): + t0 = time.perf_counter() + try: + yield + finally: + if _PROV_STATS is not None: + _PROV_STATS.record(category, time.perf_counter() - t0) + + +async def run_team( + team: Any, + task: str, + team_name: str | None = None, + source_agent_id: str | None = None, +) -> Any: + """ + Run an AutoGen team and automatically capture provenance if the plugin is active. + + This is the module-level equivalent of ``plugin.run_team()`` — it uses the + active interceptor set by ``FlowceptAutoGenPlugin.start()``, so no explicit + plugin handle is needed. Mirrors the pattern of ``openai_chat()`` and + ``anthropic_chat()`` which also use ``_ACTIVE_INTERCEPTOR`` directly. + + If the plugin is not active, the team runs normally without any provenance + overhead. + + Parameters + ---------- + team : RoundRobinGroupChat | SelectorGroupChat | any team with run_stream + The AutoGen team to run. + task : str + The task / initial message to send to the team. + team_name : str, optional + Human-readable name for this run in provenance records. + source_agent_id : str, optional + ID of an upstream agent for cross-framework linkage. + + Returns + ------- + TaskResult + The final result returned by AutoGen. + """ + interceptor = _ACTIVE_INTERCEPTOR + name = team_name or getattr(team, "name", None) or "autogen_team" + if interceptor is None: + # Plugin not active — run without provenance + from autogen_agentchat.base import TaskResult + from autogen_core import CancellationToken + + final = None + async for item in team.run_stream(task=task, cancellation_token=CancellationToken()): + if isinstance(item, TaskResult): + final = item + return final + + # Wrap pre-existing agents created before the plugin started + _ensure_agents_wrapped(team) + return await _run_with_provenance( + team=team, + task=task, + interceptor=interceptor, + stats=_PROV_STATS, + team_name=name, + source_agent_id=source_agent_id, + ) + + +def record_llm_call(payload: dict) -> None: + """ + Public API to record an LLM call into the active FlowCept provenance graph. + + Converts the payload into a TaskObject (subtype=llm_call) and routes it + through the active interceptor. No-ops if the plugin has not been started. + + Minimum payload keys: + type : "chat_completion" + model : str + text : str + usage : dict with prompt_tokens / completion_tokens / total_tokens + """ + interceptor = _ACTIVE_INTERCEPTOR + if interceptor is None: + return + import uuid as _uuid + + with _timed("record_llm_call"): + elapsed = payload.get("elapsed_s", 0.0) + now = time.time() + model = payload.get("model_used") or payload.get("model", "unknown") + + used: dict = {} + for k in ( + "model", + "model_used", + "messages", + "user_prompt", + "system_prompt", + "temperature", + "top_p", + "max_tokens", + "reasoning_effort", + "tools_provided", + "tool_choice", + "stop_sequences", + "top_k", + "thinking_budget_tokens", + ): + if k in payload: + used[k] = payload[k] + if payload.get("temperature_suppressed"): + used["temperature_suppressed"] = True + + generated: dict = {} + for k in ( + "text", + "finish_reason", + "stop_reason", + "stop_sequence", + "tool_calls", + "tool_uses", + "thinking_text", + "system_fingerprint", + "response_id", + "usage", + "elapsed_s", + ): + if k in payload: + generated[k] = payload[k] + if "error" in payload: + generated["error"] = str(payload["error"]) + + # Propagate agent linkage from contextvars (set by _run_with_provenance) + parent_task_id = _current_autogen_task_id.get(None) + agent_id = payload.get("context", {}).get("agent_name") or _current_autogen_agent_id.get(None) + + task: dict = { + "task_id": str(_uuid.uuid4()), + "subtype": "llm_call", + "activity_id": model, + "started_at": now - elapsed, + "ended_at": now, + "status": "ERROR" if "error" in payload else "FINISHED", + "used": used, + "generated": generated, + "custom_metadata": { + "model": model, + "framework": payload.get("context", {}).get("framework", ""), + "context": payload.get("context", {}), + }, + } + if parent_task_id: + task["parent_task_id"] = parent_task_id + if agent_id: + task["custom_metadata"]["agent_id"] = agent_id + interceptor.intercept_task(task) + + +# --------------------------------------------------------------------------- +# FlowceptModelClient — wraps any AutoGen ChatCompletionClient to capture +# full LLM call details (model, temperature, reasoning, usage, response, …) +# --------------------------------------------------------------------------- + + +class FlowceptModelClient: + """ + Wrap any AutoGen ``ChatCompletionClient`` for provenance capture. + + Records every LLM call as a FlowCept provenance record (subtype=llm_call) + via ``record_llm_call()``. + + Usage:: + + from autogen_ext.models.openai import OpenAIChatCompletionClient + from flowcept.agents.autogen.autogen_plugin import FlowceptModelClient + + real_client = OpenAIChatCompletionClient(model="gpt-4o-mini") + wrapped = FlowceptModelClient(real_client, agent_name="my-agent") + + agent = AssistantAgent(name="my-agent", model_client=wrapped) + + Every call to ``create()`` / ``create_stream()`` will be captured with: + - model, temperature, reasoning_effort, top_p, max_tokens + - prompt messages, system prompt + - response text, finish reason, tool calls + - prompt/completion/total token usage + - elapsed wall-clock time + """ + + def __init__(self, client: Any, agent_name: str | None = None, context: dict | None = None): + self._inner = client + self._agent_name = agent_name + self._context: dict = context or {} + + # ---- forward everything to the inner client ---- + + @property + def model_info(self): + """Return the wrapped client's model info.""" + return self._inner.model_info + + def actual_usage(self): + """Return the wrapped client's actual token usage.""" + return self._inner.actual_usage() + + def total_usage(self): + """Return the wrapped client's total token usage.""" + return self._inner.total_usage() + + def count_tokens(self, *args, **kwargs): + """Count tokens via the wrapped client.""" + return self._inner.count_tokens(*args, **kwargs) + + def remaining_tokens(self, *args, **kwargs): + """Return remaining tokens via the wrapped client.""" + return self._inner.remaining_tokens(*args, **kwargs) + + async def close(self): + """Close the wrapped client.""" + return await self._inner.close() + + # ---- instrumented create ---- + + async def create(self, messages, **kwargs) -> Any: + """Run the wrapped client's create() and record the call as provenance.""" + import time as _time + + t0 = _time.time() + result = await self._inner.create(messages, **kwargs) + elapsed = _time.time() - t0 + + self._record(messages, result, elapsed, kwargs) + return result + + async def create_stream(self, messages, **kwargs): + """Stream the wrapped client's create_stream() and record the call as provenance.""" + import time as _time + + t0 = _time.time() + final = None + async for chunk in self._inner.create_stream(messages, **kwargs): + yield chunk + final = chunk + elapsed = _time.time() - t0 + + if final is not None: + self._record(messages, final, elapsed, kwargs) + + def _record(self, messages: Any, result: Any, elapsed: float, kwargs: dict) -> None: + """Build a payload and call record_llm_call().""" + try: + # Extract model name: try direct attr first, then model_info.family fallback + mi = getattr(self._inner, "model_info", {}) or {} + model = ( + self._context.get("model_name") + or getattr(self._inner, "model", None) + or getattr(self._inner, "_model_name", None) + or getattr(mi, "model", None) + or (mi.get("family") if hasattr(mi, "get") else None) + or "unknown" + ) + + # Extract usage + usage_obj = getattr(result, "usage", None) + usage: dict = {} + if usage_obj is not None: + usage["prompt_tokens"] = getattr(usage_obj, "prompt_tokens", 0) + usage["completion_tokens"] = getattr(usage_obj, "completion_tokens", 0) + usage["total_tokens"] = usage["prompt_tokens"] + usage["completion_tokens"] + + # Extract content + content = getattr(result, "content", "") + if not isinstance(content, str): + content = str(content) + + # Extract messages list for provenance (truncated) + msgs_clip = [] + for m in messages or []: + if hasattr(m, "model_dump"): + msgs_clip.append(m.model_dump()) + elif hasattr(m, "__dict__"): + msgs_clip.append(vars(m)) + else: + msgs_clip.append(str(m)) + + ctx = dict(self._context) + if self._agent_name: + ctx["agent_name"] = self._agent_name + ctx["framework"] = "autogen" + + record_llm_call( + { + "type": "chat_completion", + "model": model, + "messages": msgs_clip, + "temperature": kwargs.get("temperature"), + "top_p": kwargs.get("top_p"), + "max_tokens": kwargs.get("max_tokens"), + "reasoning_effort": kwargs.get("reasoning_effort"), + "text": content, + "finish_reason": getattr(result, "finish_reason", None), + "usage": usage, + "elapsed_s": elapsed, + "context": ctx, + } + ) + except Exception: + pass # never break the agent run due to provenance capture + + +# --------------------------------------------------------------------------- +# FlowceptAnthropicClient — wraps anthropic.Anthropic / AsyncAnthropic to +# capture full provenance for every messages.create / messages.stream call. +# --------------------------------------------------------------------------- + + +class FlowceptAnthropicClient: + """ + Wrap an ``anthropic.Anthropic`` (or ``AsyncAnthropic``) client for provenance capture. + + Records every ``messages.create`` / ``messages.stream`` call as a FlowCept + provenance record (subtype=llm_call) via ``record_llm_call()``. + + Usage (sync):: + + import anthropic + from flowcept.agents.autogen.autogen_plugin import FlowceptAnthropicClient + + client = FlowceptAnthropicClient( + anthropic.Anthropic(), + agent_name="my-agent", + ) + response = client.messages.create( + model="claude-haiku-4-5-20251001", + max_tokens=1024, + messages=[{"role": "user", "content": "Hello"}], + ) + + Usage (async):: + + client = FlowceptAnthropicClient( + anthropic.AsyncAnthropic(), + agent_name="my-agent", + ) + response = await client.messages.create(...) + + Every call captures: model, temperature, max_tokens, top_p, top_k, + stop_sequences, thinking, messages, response text, thinking blocks, + tool uses, stop_reason, token usage, and elapsed wall-clock time. + """ + + def __init__(self, client: Any, agent_name: str | None = None, context: dict | None = None): + self._inner = client + self._agent_name = agent_name + self._context: dict = context or {} + self.messages = _FlowceptAnthropicMessages(client.messages, agent_name, self._context) + + def __getattr__(self, name: str) -> Any: + """Delegate attribute access to the wrapped client.""" + return getattr(self._inner, name) + + +class _FlowceptAnthropicMessages: + """Proxy for client.messages that intercepts create() and stream().""" + + def __init__(self, messages_resource: Any, agent_name: str | None, context: dict): + self._inner = messages_resource + self._agent_name = agent_name + self._context = context + + def __getattr__(self, name: str) -> Any: + return getattr(self._inner, name) + + def create(self, **kwargs) -> Any: + import time as _time + + t0 = _time.time() + result = self._inner.create(**kwargs) + elapsed = _time.time() - t0 + self._record(kwargs, result, elapsed, is_stream=False) + return result + + async def async_create(self, **kwargs) -> Any: + import time as _time + + t0 = _time.time() + result = await self._inner.create(**kwargs) + elapsed = _time.time() - t0 + self._record(kwargs, result, elapsed, is_stream=False) + return result + + def stream(self, **kwargs): + import time as _time + + t0 = _time.time() + ctx_mgr = self._inner.stream(**kwargs) + # Wrap the context manager to capture timing on exit + return _FlowceptAnthropicStream(ctx_mgr, kwargs, self._record, t0) + + def _record(self, kwargs: dict, result: Any, elapsed: float, is_stream: bool = False) -> None: + try: + model = kwargs.get("model", "unknown") + + # Extract response fields + text = "" + thinking_text = "" + tool_uses = [] + for block in getattr(result, "content", []): + btype = getattr(block, "type", None) + if btype == "text": + text += getattr(block, "text", "") + elif btype == "thinking": + thinking_text += getattr(block, "thinking", "") + elif btype == "tool_use": + tool_uses.append( + { + "id": getattr(block, "id", None), + "name": getattr(block, "name", None), + "input": getattr(block, "input", None), + } + ) + + usage = getattr(result, "usage", None) + usage_dict = {} + if usage is not None: + for attr in ("input_tokens", "output_tokens", "cache_creation_input_tokens", "cache_read_input_tokens"): + val = getattr(usage, attr, None) + if val is not None: + usage_dict[attr] = val + + ctx = dict(self._context) + if self._agent_name: + ctx["agent_name"] = self._agent_name + ctx["framework"] = "autogen" + + payload: dict = { + "type": "chat_completion", + "model": model, + "model_used": getattr(result, "model", model), + "messages": kwargs.get("messages"), + "system_prompt": kwargs.get("system"), + "max_tokens": kwargs.get("max_tokens"), + "temperature": kwargs.get("temperature"), + "top_p": kwargs.get("top_p"), + "top_k": kwargs.get("top_k"), + "stop_sequences": kwargs.get("stop_sequences"), + "thinking_budget_tokens": (kwargs.get("thinking") or {}).get("budget_tokens"), + "tools_provided": [t.get("name") for t in (kwargs.get("tools") or [])], + "tool_choice": kwargs.get("tool_choice"), + "text": text, + "thinking_text": thinking_text or None, + "finish_reason": getattr(result, "stop_reason", None), + "stop_sequence": getattr(result, "stop_sequence", None), + "tool_uses": tool_uses or None, + "response_id": getattr(result, "id", None), + "usage": usage_dict, + "elapsed_s": elapsed, + "context": ctx, + } + record_llm_call({k: v for k, v in payload.items() if v is not None}) + except Exception: + pass # never break user code due to provenance capture + + +class _FlowceptAnthropicStream: + """Thin wrapper around anthropic streaming context manager.""" + + def __init__(self, ctx_mgr: Any, kwargs: dict, record_fn, t0: float): + self._ctx_mgr = ctx_mgr + self._kwargs = kwargs + self._record_fn = record_fn + self._t0 = t0 + self._stream = None + + def __enter__(self): + self._stream = self._ctx_mgr.__enter__() + return self._stream + + def __exit__(self, *args): + import time as _time + + result = None + try: + result = self._stream.get_final_message() + except Exception: + pass + elapsed = _time.time() - self._t0 + if result is not None: + self._record_fn(self._kwargs, result, elapsed, is_stream=True) + return self._ctx_mgr.__exit__(*args) + + +def openai_chat( + prompt: str, + model: str = "gpt-4o-mini", + system: str = "You are a helpful assistant.", + temperature: float | None = 0.3, + top_p: float | None = None, + max_tokens: int | None = None, + n: int = 1, + stop: list[str] | str | None = None, + frequency_penalty: float = 0.0, + presence_penalty: float = 0.0, + seed: int | None = None, + reasoning_effort: str | None = None, + response_format: dict | None = None, + tools: list | None = None, + tool_choice: str | dict | None = None, + user: str | None = None, + context: dict | None = None, +) -> str: + """ + Make an OpenAI chat completion call and record it for FlowCept provenance. + + Captures all request parameters and response fields as a child TaskObject. + No-ops gracefully if OPENAI_API_KEY is not set or the plugin is not started. + """ + import openai as _openai + import os as _os + import time as _time + + client = _openai.OpenAI(api_key=_os.environ.get("OPENAI_API_KEY")) + messages = [ + {"role": "system", "content": system}, + {"role": "user", "content": prompt}, + ] + req: dict = {"model": model, "messages": messages, "n": n} + if temperature is not None: + req["temperature"] = temperature + if top_p is not None: + req["top_p"] = top_p + if max_tokens is not None: + req["max_completion_tokens"] = max_tokens + if stop is not None: + req["stop"] = stop + if frequency_penalty != 0.0: + req["frequency_penalty"] = frequency_penalty + if presence_penalty != 0.0: + req["presence_penalty"] = presence_penalty + if seed is not None: + req["seed"] = seed + if reasoning_effort is not None: + req["reasoning_effort"] = reasoning_effort + if response_format is not None: + req["response_format"] = response_format + if tools is not None: + req["tools"] = tools + if tool_choice is not None: + req["tool_choice"] = tool_choice + if user is not None: + req["user"] = user + + t0 = _time.time() + response = client.chat.completions.create(**req) + elapsed = _time.time() - t0 + + choice = response.choices[0] + text = choice.message.content or "" + usage = response.usage or {} + + usage_dict: dict = {} + for attr in ("prompt_tokens", "completion_tokens", "total_tokens"): + if hasattr(usage, attr): + usage_dict[attr] = getattr(usage, attr) + if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details: + ctd = usage.completion_tokens_details + usage_dict["reasoning_tokens"] = getattr(ctd, "reasoning_tokens", None) + usage_dict["accepted_prediction_tokens"] = getattr(ctd, "accepted_prediction_tokens", None) + usage_dict["rejected_prediction_tokens"] = getattr(ctd, "rejected_prediction_tokens", None) + if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details: + ptd = usage.prompt_tokens_details + usage_dict["cached_tokens"] = getattr(ptd, "cached_tokens", None) + + tool_calls = None + if choice.message.tool_calls: + tool_calls = [ + { + "id": tc.id, + "type": tc.type, + "function": {"name": tc.function.name, "arguments": tc.function.arguments}, + } + for tc in choice.message.tool_calls + ] + + record_llm_call( + { + "type": "chat_completion", + "model": model, + "model_used": response.model, + "messages": messages, + "user_prompt": prompt, + "system_prompt": system, + "temperature": temperature, + "temperature_suppressed": temperature is None, + "top_p": top_p, + "max_tokens": max_tokens, + "n": n, + "stop": stop, + "frequency_penalty": frequency_penalty, + "presence_penalty": presence_penalty, + "seed": seed, + "reasoning_effort": reasoning_effort, + "response_format": response_format, + "tools_provided": [t.get("function", {}).get("name") for t in (tools or [])], + "tool_choice": tool_choice, + "text": text, + "finish_reason": choice.finish_reason, + "tool_calls": tool_calls, + "system_fingerprint": getattr(response, "system_fingerprint", None), + "response_id": response.id, + "created": response.created, + "usage": usage_dict, + "elapsed_s": elapsed, + "context": context or {}, + } + ) + return text + + +def anthropic_chat( + prompt: str, + model: str = "claude-haiku-4-5-20251001", + system: str = "You are a helpful assistant.", + max_tokens: int = 1024, + temperature: float | None = 1.0, + top_p: float | None = None, + top_k: int | None = None, + stop_sequences: list[str] | None = None, + tools: list | None = None, + tool_choice: dict | None = None, + thinking: dict | None = None, + metadata: dict | None = None, + context: dict | None = None, +) -> str: + """ + Make an Anthropic (Claude) chat completion call and record it for FlowCept provenance. + + Captures all request/response fields — including thinking blocks, tool use, + and cache token counts — as a child TaskObject. + """ + import anthropic as _anthropic + import os as _os + import time as _time + + client = _anthropic.Anthropic(api_key=_os.environ.get("ANTHROPIC_API_KEY")) + req: dict = { + "model": model, + "max_tokens": max_tokens, + "system": system, + "messages": [{"role": "user", "content": prompt}], + } + if temperature is not None: + req["temperature"] = temperature + if top_p is not None: + req["top_p"] = top_p + if top_k is not None: + req["top_k"] = top_k + if stop_sequences: + req["stop_sequences"] = stop_sequences + if tools: + req["tools"] = tools + if tool_choice: + req["tool_choice"] = tool_choice + if thinking: + req["thinking"] = thinking + if metadata: + req["metadata"] = metadata + + t0 = _time.time() + response = client.messages.create(**req) + elapsed = _time.time() - t0 + + text = "" + thinking_text = "" + tool_uses = [] + for block in response.content: + if block.type == "text": + text += block.text + elif block.type == "thinking": + thinking_text += getattr(block, "thinking", "") + elif block.type == "tool_use": + tool_uses.append({"id": block.id, "name": block.name, "input": block.input}) + + usage = response.usage + usage_dict = { + "input_tokens": getattr(usage, "input_tokens", None), + "output_tokens": getattr(usage, "output_tokens", None), + "cache_creation_input_tokens": getattr(usage, "cache_creation_input_tokens", None), + "cache_read_input_tokens": getattr(usage, "cache_read_input_tokens", None), + } + + record_llm_call( + { + "type": "chat_completion", + "model": model, + "model_used": response.model, + "messages": req["messages"], + "user_prompt": prompt, + "system_prompt": system, + "max_tokens": max_tokens, + "temperature": temperature, + "temperature_suppressed": temperature is None, + "top_p": top_p, + "top_k": top_k, + "stop_sequences": stop_sequences, + "thinking_budget_tokens": (thinking or {}).get("budget_tokens"), + "tools_provided": [t.get("name") for t in (tools or [])], + "tool_choice": tool_choice, + "text": text, + "thinking_text": thinking_text if thinking_text else None, + "finish_reason": response.stop_reason, + "stop_sequence": response.stop_sequence, + "tool_uses": tool_uses if tool_uses else None, + "response_id": response.id, + "usage": usage_dict, + "elapsed_s": elapsed, + "context": context or {}, + } + ) + return text + + +# --------------------------------------------------------------------------- +# Public plugin class +# --------------------------------------------------------------------------- + + +class FlowceptAutoGenPlugin: + """ + FlowCept provenance plugin for AutoGen (autogen_agentchat 0.7+) team runs. + + Wraps ``team.run_stream()`` to capture provenance without patching the + AutoGen library. Each message yielded by the stream becomes a TaskObject + child of the overall run TaskObject. + + Parameters + ---------- + config : dict, optional + Plugin configuration keys: + enabled (bool, default True) + workflow_name (str, default "autogen-workflow") + performance_tracking (bool, default True) + perf_csv (str, optional) — explicit path for timing CSV. + + Usage + ----- + Standalone:: + + plugin = FlowceptAutoGenPlugin(config={"workflow_name": "my-team"}) + plugin.start() + result = asyncio.run(plugin.run_team(team, "Solve the problem")) + plugin.stop() + + Shared with Academy plugin:: + + academy_plugin = FlowceptAcademyPlugin(config={...}).start() + autogen_plugin = FlowceptAutoGenPlugin.from_academy_plugin(academy_plugin) + result = asyncio.run(autogen_plugin.run_team(team, "task")) + academy_plugin.stop() + """ + + def __init__(self, config: dict | None = None, _shared_interceptor=None) -> None: + cfg = config or {} + self._enabled: bool = cfg.get("enabled", True) + self._workflow_name: str = cfg.get("workflow_name", "autogen-workflow") + self._campaign_id: str | None = cfg.get("campaign_id", None) + self._perf_tracking: bool = cfg.get("performance_tracking", True) + self._perf_csv: str | None = cfg.get("perf_csv", None) + self._shared_interceptor = _shared_interceptor + self._interceptor = _shared_interceptor or _AutoGenInterceptor() + self._owns_interceptor: bool = _shared_interceptor is None + self._stats: _ProvenanceStats | None = None + self._started = False + + @classmethod + def from_academy_plugin( + cls, + academy_plugin: Any, + config: dict | None = None, + ) -> "FlowceptAutoGenPlugin": + """Create an AutoGen plugin that shares the buffer of a running FlowceptAcademyPlugin.""" + interceptor = academy_plugin._interceptor + inst = cls(config=config, _shared_interceptor=interceptor) + inst._started = True + global _ACTIVE_INTERCEPTOR + _ACTIVE_INTERCEPTOR = interceptor + + inst._stats = _ProvenanceStats() if (config or {}).get("performance_tracking", True) else None + global _PROV_STATS + _PROV_STATS = inst._stats + return inst + + def start(self) -> "FlowceptAutoGenPlugin": + """Start provenance capture, initializing the interceptor if this plugin owns it.""" + if not self._enabled or self._started: + return self + if not self._owns_interceptor: + return self + try: + self._stats = _ProvenanceStats() if self._perf_tracking else None + global _PROV_STATS + _PROV_STATS = self._stats + self._interceptor.start(self._workflow_name, campaign_id=self._campaign_id) + self._campaign_id = self._interceptor._campaign_id + self._started = True + global _ACTIVE_INTERCEPTOR + _ACTIVE_INTERCEPTOR = self._interceptor + _install_autogen_patches() + wf_id = self._interceptor._workflow_id + print( + f"[FlowceptAutoGenPlugin] Started\n" + f" workflow_id : {wf_id}\n" + f" campaign_id : {self._interceptor._campaign_id}\n" + f" Capturing : team run (sub-workflow), messages (child tasks).", + flush=True, + ) + except Exception as e: + print( + f"[FlowceptAutoGenPlugin] WARNING: failed to start — {e!r}. Continuing without provenance capture.", + flush=True, + ) + _log.exception("FlowceptAutoGenPlugin start failed") + self._enabled = False + return self + + async def run_team( + self, + team: Any, + task: str, + team_name: str | None = None, + source_agent_id: str | None = None, + ) -> Any: + """ + Run a team and capture provenance for the entire conversation. + + Parameters + ---------- + team : RoundRobinGroupChat | SelectorGroupChat | any team with run_stream + The AutoGen team to run. + task : str + The task / initial message to send to the team. + team_name : str, optional + Human-readable name for this run in provenance records. + Defaults to the team's ``name`` attribute or "autogen_team". + source_agent_id : str, optional + ID of an upstream agent (e.g. Academy AgentId) that produced the + input data. Stored in custom_metadata for cross-framework linkage. + + Returns + ------- + TaskResult + The final result returned by AutoGen. + """ + name = team_name or getattr(team, "name", None) or "autogen_team" + if not self._started: + # plugin not started — run the team without provenance + from autogen_agentchat.base import TaskResult + from autogen_core import CancellationToken + + results = [] + async for item in team.run_stream(task=task, cancellation_token=CancellationToken()): + if isinstance(item, TaskResult): + results.append(item) + return results[-1] if results else None + + return await _run_with_provenance( + team=team, + task=task, + interceptor=self._interceptor, + stats=self._stats, + team_name=name, + source_agent_id=source_agent_id, + ) + + def stop(self) -> None: + """Stop provenance capture and flush records if this plugin owns the interceptor.""" + if not self._started: + return + if not self._owns_interceptor: + self._started = False + print( + "[FlowceptAutoGenPlugin] Detached from shared buffer (flushed by the owning plugin).", + flush=True, + ) + self._maybe_write_perf_csv() + global _PROV_STATS + _PROV_STATS = None + return + try: + _t0 = time.perf_counter() + self._interceptor.stop() + if self._stats is not None: + self._stats.record("flush", time.perf_counter() - _t0) + except Exception as e: + print(f"[FlowceptAutoGenPlugin] Warning during stop: {e!r}", flush=True) + self._started = False + _uninstall_autogen_patches() + print("[FlowceptAutoGenPlugin] Stopped.", flush=True) + self._maybe_write_perf_csv() + global _ACTIVE_INTERCEPTOR + _ACTIVE_INTERCEPTOR = None + _PROV_STATS = None + + def _maybe_write_perf_csv(self) -> None: + if self._stats is None: + return + print( + "\n[FlowceptAutoGenPlugin] Provenance overhead report:\n" + + self._stats.summary() + + "\n (N = event count; Total/Mean/Min/Max in ms/µs respectively)\n", + flush=True, + ) + wf_id = self._interceptor._workflow_id + csv_path = self._perf_csv or f"autogen_provenance_perf_{wf_id}.csv" + try: + self._stats.to_csv(csv_path, workflow_id=wf_id) + print(f"[FlowceptAutoGenPlugin] Performance stats → {csv_path}", flush=True) + except Exception as e: + print(f"[FlowceptAutoGenPlugin] Warning: could not write perf CSV — {e!r}", flush=True) + + def __enter__(self) -> "FlowceptAutoGenPlugin": + """Start the plugin when entering a context manager block.""" + return self.start() + + def __exit__(self, *_: Any) -> None: + """Stop the plugin when exiting a context manager block.""" + self.stop() diff --git a/src/flowcept/agents/chat_orchestration/tool_registry.py b/src/flowcept/agents/chat_orchestration/tool_registry.py index 7a958831..e9c90f19 100644 --- a/src/flowcept/agents/chat_orchestration/tool_registry.py +++ b/src/flowcept/agents/chat_orchestration/tool_registry.py @@ -16,6 +16,7 @@ from pydantic import create_model from flowcept.agents.mcp.mcp_client import run_tool +from flowcept.agents.mcp.mcp_tools import analysis_mcp_tools as _analysis from flowcept.agents.mcp.mcp_tools import db_query_mcp_tools as _db from flowcept.agents.mcp.mcp_tools import df_query_mcp_tools as _df from flowcept.agents.mcp.mcp_tools import dashboard_mcp_tools as _dash @@ -273,6 +274,105 @@ def update_dashboard( return _run_mcp(_dash.db_update_dashboard.__name__, dashboard_id=dashboard_id, spec=spec or {}) +def summarize_execution( + tool_context: str, + context: Optional[Dict[str, Any]], + workflow_id: Optional[str] = None, +) -> str: + """Summarize an execution: task counts by activity/subtype, statuses, durations, token usage. + + DB mode: analyzes DB records, optionally scoped by ``workflow_id``. + DF mode: analyzes the records loaded in the agent's in-memory context. + """ + effective_id = workflow_id or (context or {}).get("workflow_id") + if tool_context == "df": + return _run_mcp(_analysis.df_summarize_execution.__name__, workflow_id=effective_id) + return _run_mcp(_analysis.db_summarize_execution.__name__, workflow_id=effective_id) + + +def analyze_errors( + tool_context: str, + context: Optional[Dict[str, Any]], + workflow_id: Optional[str] = None, +) -> str: + """Analyze failed tasks: per-activity error rates, stderr excerpts, first/last failure times. + + DB mode: analyzes DB records, optionally scoped by ``workflow_id``. + DF mode: analyzes the records loaded in the agent's in-memory context. + """ + if tool_context == "df": + return _run_mcp(_analysis.df_analyze_errors.__name__) + effective_id = workflow_id or (context or {}).get("workflow_id") + return _run_mcp(_analysis.db_analyze_errors.__name__, workflow_id=effective_id) + + +def analyze_agent_behavior( + tool_context: str, + context: Optional[Dict[str, Any]], + workflow_id: Optional[str] = None, +) -> str: + """Profile per-agent behavior: turns, tool calls by tool, LLM calls, token usage, durations. + + DB mode: analyzes DB records, optionally scoped by ``workflow_id``. + DF mode: analyzes the records loaded in the agent's in-memory context. + """ + if tool_context == "df": + return _run_mcp(_analysis.df_agent_behavior.__name__) + effective_id = workflow_id or (context or {}).get("workflow_id") + return _run_mcp(_analysis.db_agent_behavior.__name__, workflow_id=effective_id) + + +def find_slowest_tasks( + tool_context: str, + context: Optional[Dict[str, Any]], + limit: int = 10, + workflow_id: Optional[str] = None, +) -> str: + """Find the slowest tasks, longest elapsed first, with status and parent-chain depth. + + DB mode: analyzes DB records, optionally scoped by ``workflow_id``. + DF mode: analyzes the records loaded in the agent's in-memory context. + """ + if tool_context == "df": + return _run_mcp(_analysis.df_find_slowest.__name__, limit=limit) + effective_id = workflow_id or (context or {}).get("workflow_id") + return _run_mcp(_analysis.db_find_slowest.__name__, workflow_id=effective_id, limit=limit) + + +def cross_framework_links( + tool_context: str, + context: Optional[Dict[str, Any]], + workflow_id: Optional[str] = None, +) -> str: + """List cross-framework provenance links (source_agent_id edges) and unlinked task counts. + + DB mode: analyzes DB records, optionally scoped by ``workflow_id``. + DF mode: analyzes the records loaded in the agent's in-memory context. + """ + if tool_context == "df": + return _run_mcp(_analysis.df_cross_framework_links.__name__) + effective_id = workflow_id or (context or {}).get("workflow_id") + return _run_mcp(_analysis.db_cross_framework_links.__name__, workflow_id=effective_id) + + +def compare_executions( + tool_context: str, + context: Optional[Dict[str, Any]], + workflow_id_a: str = "", + workflow_id_b: str = "", +) -> str: + """Compare two workflow executions per activity: count, duration, and error-rate deltas. + + Pass the two workflow ids to compare. Records come from the agent context + when loaded there, otherwise from the database. + """ + return _run_mcp( + _analysis.compare_executions.__name__, + workflow_id_a=workflow_id_a, + workflow_id_b=workflow_id_b, + ) + + # --------------------------------------------------------------------------- # LangChain tool builder # --------------------------------------------------------------------------- @@ -343,6 +443,12 @@ def _run(**kwargs): make_chart, highlight_lineage, fix_query, + summarize_execution, + analyze_errors, + analyze_agent_behavior, + find_slowest_tasks, + cross_framework_links, + compare_executions, ] _DASHBOARD_TOOLS = [get_dashboard, update_dashboard] diff --git a/src/flowcept/agents/claude_agent_sdk/claude_agent_sdk_plugin.py b/src/flowcept/agents/claude_agent_sdk/claude_agent_sdk_plugin.py new file mode 100644 index 00000000..3bfd1087 --- /dev/null +++ b/src/flowcept/agents/claude_agent_sdk/claude_agent_sdk_plugin.py @@ -0,0 +1,329 @@ +"""Claude Agent SDK wrapper. + +The SDK hands you an async stream of message objects. Everything provenance +needs is already in that stream — tool uses, tool results, token usage, the +final answer — so capture is a matter of reading it as it goes past rather than +instrumenting anything. + + from flowcept.agents.claude_agent_sdk.claude_agent_sdk_plugin import trace_query + + async for message in trace_query(prompt="fix the failing test"): + print(message) + +:func:`trace_query` is a drop-in for ``claude_agent_sdk.query``: same +arguments, same yielded messages, provenance as a side effect. To capture a +``ClaudeSDKClient`` conversation instead, drive :class:`ClaudeAgentTracer` +directly with the messages you receive. + +Messages are matched structurally, not with ``isinstance``. The SDK's block +classes have moved between versions and this module must import without the +SDK present at all, so it reads the shape it needs and ignores the rest. +""" + +from __future__ import annotations + +import json +from typing import Any + +from flowcept.agents.harness.config import Config +from flowcept.agents.harness.tracer import SessionTracer + +#: The tool the assistant uses to spawn a subagent. It gets a nested workflow +#: in addition to its tool record, so the subagent's work is attributable to it. +SUBAGENT_TOOL = "Task" + + +class ClaudeAgentTracer: + """Turns a Claude Agent SDK message stream into Flowcept provenance. + + Feed it every message the SDK yields, in order, then call :meth:`close`. + One turn is opened per user prompt and closed by the ``ResultMessage``. + """ + + def __init__( + self, + session_id: str | None = None, + *, + config: Config | None = None, + model: str | None = None, + harness: str = "claude_agent_sdk", + prompt: str | None = None, + tracer: SessionTracer | None = None, + ): + self.tracer = tracer or SessionTracer(harness, session_id, config=config, model=model) + self._open_turn = False + self._pending_prompt = prompt + #: tool_use_id -> tool name, so a result can name the tool it closes. + self._tools: dict[str, str] = {} + #: tool_use_id of Task calls, which also own a subagent workflow. + self._subagents: dict[str, str] = {} + self._text: list[str] = [] + #: Whether any record has been written yet. Until it has, the session + #: id is still negotiable; see :meth:`_adopt_session_id`. + self._emitted = False + self._explicit_session_id = session_id is not None + + def _adopt_session_id(self, session_id: Any) -> None: + """Take the SDK's session id, if it is not too late to. + + All provenance ids derive from the session id, so adopting the SDK's + makes a resumed conversation continue the same workflow instead of + forking a new one. Once a record has been written the id is load- + bearing and changing it would orphan everything already emitted, so + this only ever fires before the first one. + """ + if self._emitted or self._explicit_session_id or not isinstance(session_id, str): + return + self.tracer.session_id = session_id + + def _ensure_started(self) -> None: + if self._emitted: + return + self._emitted = True + self.tracer.start() + if self._pending_prompt is not None: + self.begin_turn(self._pending_prompt) + + # -- turns --------------------------------------------------------------- + + def begin_turn(self, prompt: str | None = None) -> None: + """Open a turn for a user prompt.""" + self._emitted = True + self._text = [] + self._open_turn = True + self._pending_prompt = None + self.tracer.prompt(prompt) + + # -- the stream ---------------------------------------------------------- + + def handle(self, message: Any) -> None: + """Record one streamed message. Unknown messages are ignored.""" + kind = _message_kind(message) + if kind == "system": + # Handled first and without starting: the SDK's init message is + # where the real session id arrives, ahead of everything else. + self._on_system(message) + return + self._adopt_session_id(getattr(message, "session_id", None)) + self._ensure_started() + if kind == "assistant": + self._on_assistant(message) + elif kind == "user": + self._on_user(message) + elif kind == "result": + self._on_result(message) + + def _on_assistant(self, message: Any) -> None: + model = getattr(message, "model", None) + if model: + self.tracer.model = model + if not self._open_turn: + self.begin_turn(self._pending_prompt) + + for block in _blocks(message): + text = getattr(block, "text", None) + if isinstance(text, str): + self._text.append(text) + continue + + name = getattr(block, "name", None) + tool_use_id = getattr(block, "id", None) + if not (name and tool_use_id): + continue + + tool_input = getattr(block, "input", None) + self._tools[tool_use_id] = name + self.tracer.tool_start(name, tool_input, tool_use_id=tool_use_id) + + if name == SUBAGENT_TOOL: + arguments = tool_input if isinstance(tool_input, dict) else {} + self._subagents[tool_use_id] = self.tracer.subagent_start( + arguments.get("subagent_type") or arguments.get("description") or "subagent", + agent_ref=tool_use_id, + prompt=arguments.get("prompt"), + ) + + def _on_user(self, message: Any) -> None: + """A user message in the stream is the harness returning tool results.""" + for block in _blocks(message): + tool_use_id = getattr(block, "tool_use_id", None) + if not tool_use_id: + continue + content = getattr(block, "content", None) + is_error = bool(getattr(block, "is_error", False)) + self.tracer.tool_end( + tool_use_id, + name=self._tools.pop(tool_use_id, None), + tool_response=None if is_error else content, + error=_as_text(content) if is_error else None, + ) + if tool_use_id in self._subagents: + self.tracer.subagent_stop( + self._subagents.pop(tool_use_id), + response=None if is_error else _as_text(content), + error=_as_text(content) if is_error else None, + ) + + def _on_result(self, message: Any) -> None: + usage = _as_dict(getattr(message, "usage", None)) + cost = getattr(message, "total_cost_usd", None) + if usage is not None and cost is not None: + usage = {**usage, "total_cost_usd": cost} + + response = getattr(message, "result", None) + if not isinstance(response, str): + response = "".join(self._text) or None + + error = None + if getattr(message, "is_error", False): + error = _as_text(response) or "the SDK reported an error result" + + self.tracer.turn_end(response=response, usage=usage, error=error) + self._open_turn = False + self._text = [] + + def _on_system(self, message: Any) -> None: + subtype = getattr(message, "subtype", None) + data = getattr(message, "data", None) + if isinstance(data, dict): + self._adopt_session_id(data.get("session_id")) + model = data.get("model") + if isinstance(model, str): + self.tracer.model = model + self._adopt_session_id(getattr(message, "session_id", None)) + + if subtype == "init": + # Nothing happened yet worth a record; the id and model were the + # point of this message. + return + self._ensure_started() + if subtype == "compact_boundary": + self.tracer.compact(source=subtype) + elif subtype: + self.tracer.notify(subtype) + + # -- teardown ------------------------------------------------------------ + + def close(self, *, error: str | None = None) -> None: + """Close any open tool calls, the turn, and the session. + + A run that produced nothing at all is left unrecorded: an empty session + is noise, not provenance. A run that had a prompt still gets one, since + "we asked and got nothing back" is worth knowing. + """ + if not self._emitted: + if self._pending_prompt is None and error is None: + return + self._ensure_started() + for tool_use_id, name in list(self._tools.items()): + self.tracer.tool_end(tool_use_id, name=name, error="never returned a result") + self._tools.clear() + for ref in list(self._subagents.values()): + self.tracer.subagent_stop(ref, error="never returned a result") + self._subagents.clear() + if self._open_turn: + self.tracer.turn_end(response="".join(self._text) or None, error=error) + self._open_turn = False + self.tracer.end(source="error" if error else "completed") + + def __enter__(self) -> ClaudeAgentTracer: + return self + + def __exit__(self, exc_type, exc, tb) -> bool: + self.close(error=f"{exc_type.__name__}: {exc}" if exc_type else None) + return False + + +async def trace_query(prompt: Any, options: Any = None, *, config: Config | None = None, **kwargs: Any): + """``claude_agent_sdk.query`` with provenance capture. + + Yields exactly what ``query`` yields; the capture is a side effect, so this + can be substituted for the original call without changing the consumer. + """ + from claude_agent_sdk import query # imported here: the SDK is an extra + + tracer = ClaudeAgentTracer( + config=config, + prompt=prompt if isinstance(prompt, str) else None, + ) + try: + async for message in query(prompt=prompt, options=options, **kwargs): + try: + tracer.handle(message) + except Exception: + # Capture must never break the stream it is observing. + pass + yield message + except BaseException as exc: + tracer.close(error=f"{type(exc).__name__}: {exc}") + raise + else: + tracer.close() + + +# -- structural message matching --------------------------------------------- + + +def _message_kind(message: Any) -> str | None: + """Classify a message by class name, falling back to its shape.""" + name = type(message).__name__ + for candidate in ("Assistant", "User", "Result", "System"): + if name.startswith(candidate): + return candidate.lower() + + if hasattr(message, "num_turns") or hasattr(message, "total_cost_usd"): + return "result" + if hasattr(message, "data") and hasattr(message, "subtype"): + return "system" + if hasattr(message, "content"): + # Only the assistant's messages name the model that produced them. + return "assistant" if getattr(message, "model", None) else "user" + return None + + +def _blocks(message: Any) -> list[Any]: + content = getattr(message, "content", None) + if isinstance(content, list): + return content + if isinstance(content, str): + return [_TextBlock(content)] + return [] + + +class _TextBlock: + """Wraps bare string content so callers see one uniform block shape.""" + + __slots__ = ("text",) + + def __init__(self, text: str): + self.text = text + + +def _as_dict(value: Any) -> dict[str, Any] | None: + if value is None or isinstance(value, dict): + return value + for attribute in ("model_dump", "to_dict", "_asdict"): + method = getattr(value, attribute, None) + if callable(method): + try: + result = method() + except Exception: + continue + if isinstance(result, dict): + return result + return getattr(value, "__dict__", None) or None + + +def _as_text(value: Any) -> str | None: + if value is None or isinstance(value, str): + return value + if isinstance(value, list): + # Tool results arrive as content blocks; keep the text, drop the rest. + parts = [b.get("text") if isinstance(b, dict) else getattr(b, "text", None) for b in value] + joined = "".join(p for p in parts if isinstance(p, str)) + if joined: + return joined + try: + return json.dumps(value, default=repr) + except (TypeError, ValueError): + return repr(value) diff --git a/src/flowcept/agents/claude_code/claude_code_plugin.py b/src/flowcept/agents/claude_code/claude_code_plugin.py new file mode 100644 index 00000000..529aed91 --- /dev/null +++ b/src/flowcept/agents/claude_code/claude_code_plugin.py @@ -0,0 +1,198 @@ +"""Claude Code adapter. + +Claude Code delivers each lifecycle event as a JSON object on stdin to a hook +command, one process per event. This module is that command: it maps the event +onto a :class:`~flowcept.agents.harness.events.HarnessEvent` and hands it to the +recorder. + +Two things shape the implementation: + +*Field names are read defensively.* The hook payload is versioned with the +CLI, and only a subset of it is contractually stable across versions. Known +names are tried in order and the untouched payload is kept under +``custom_metadata.raw_event``, so a renamed field degrades the mapping rather +than losing the event. + +*Nothing is ever written to stdout.* On ``UserPromptSubmit``, ``SessionStart``, +and ``UserPromptExpansion`` a hook's stdout is injected into the model's +context. A provenance hook that printed anything would silently edit the +conversation it is supposed to be observing. +""" + +from __future__ import annotations + +import sys +from typing import Any + +from flowcept.agents.harness.config import Config +from flowcept.agents.harness.events import HarnessEvent +from flowcept.agents.harness.recorder import Recorder +from flowcept.agents.harness.runtime import log_error, project_dir_from_env, read_stdin_json, run_capture +from flowcept.agents.harness.vocab import EventKind + +HARNESS = "claude_code" + +#: Claude Code hook event -> normalized event kind. Events not listed here are +#: intentionally ignored: they carry no provenance we do not already capture, +#: and every extra hook is latency on the interactive path. +EVENT_MAP: dict[str, str] = { + "SessionStart": EventKind.SESSION_START, + "SessionEnd": EventKind.SESSION_END, + "UserPromptSubmit": EventKind.PROMPT, + "Stop": EventKind.TURN_END, + "StopFailure": EventKind.TURN_END, + "PreToolUse": EventKind.TOOL_PRE, + "PostToolUse": EventKind.TOOL_POST, + "PostToolUseFailure": EventKind.TOOL_ERROR, + "SubagentStart": EventKind.SUBAGENT_START, + "SubagentStop": EventKind.SUBAGENT_STOP, + "Notification": EventKind.NOTIFICATION, + "PreCompact": EventKind.COMPACT, + "PostCompact": EventKind.COMPACT, +} + + +def _first(payload: dict[str, Any], *names: str) -> Any: + """Return the first present, non-None value among ``names``.""" + for name in names: + if name in payload and payload[name] is not None: + return payload[name] + return None + + +def _text(value: Any) -> str | None: + """Coerce a message-ish value to text. + + ``last_assistant_message`` is usually a string but can arrive as the + content-block list the API uses. + """ + if value is None or isinstance(value, str): + return value + if isinstance(value, list): + parts = [] + for block in value: + if isinstance(block, str): + parts.append(block) + elif isinstance(block, dict): + text = block.get("text") or block.get("content") + if isinstance(text, str): + parts.append(text) + return "\n".join(parts) if parts else None + if isinstance(value, dict): + text = value.get("text") or value.get("content") or value.get("message") + return text if isinstance(text, str) else None + return str(value) + + +def _effort(payload: dict[str, Any]) -> str | None: + effort = payload.get("effort") + if isinstance(effort, dict): + level = effort.get("level") + return level if isinstance(level, str) else None + return effort if isinstance(effort, str) else None + + +def to_event(payload: dict[str, Any]) -> HarnessEvent | None: + """Map a Claude Code hook payload onto a normalized event.""" + hook_name = payload.get("hook_event_name") + kind = EVENT_MAP.get(hook_name or "") + if kind is None: + return None + + session_id = payload.get("session_id") + if not session_id: + # Without a session id nothing can be correlated; drop rather than + # invent an id that would fragment the workflow across processes. + return None + + cwd = payload.get("cwd") + + # `agent_id` is only present when the hook fires inside a subagent. For + # subagent lifecycle events it identifies the subagent being started or + # stopped; for tool events it says which subagent owns the call. + agent_ref = payload.get("agent_id") + agent_name = payload.get("agent_type") + + # `flowcept_source_agent_id` is Flowcept's own extension key: something + # that injects payload fields (a wrapper, a test, a future harness knob) + # can name a framework-emitted task/agent id to link back to. It wins over + # the FLOWCEPT_HARNESS_SOURCE_AGENT_ID env fallback read by the recorder. + source_agent_id = payload.get("flowcept_source_agent_id") + + error = None + if hook_name == "PostToolUseFailure": + error = _text(_first(payload, "error", "tool_error", "tool_response")) or "tool failed" + elif hook_name == "StopFailure": + error = _text(_first(payload, "error", "reason", "error_type")) or "turn failed" + + event = HarnessEvent( + kind=kind, + harness=HARNESS, + session_id=str(session_id), + cwd=cwd, + project_dir=project_dir_from_env(cwd), + model=_first(payload, "model", "model_id"), + permission_mode=payload.get("permission_mode"), + effort=_effort(payload), + source=_text(_first(payload, "source", "reason", "trigger", "notification_type")), + prompt_id=payload.get("prompt_id"), + prompt=_text(_first(payload, "prompt", "user_prompt")), + response=_text(_first(payload, "last_assistant_message", "response")), + tool_name=payload.get("tool_name"), + tool_use_id=payload.get("tool_use_id"), + tool_input=_first(payload, "tool_input", "tool_args"), + tool_response=_first(payload, "tool_response", "tool_result", "tool_output"), + error=error, + agent_name=agent_name, + agent_ref=agent_ref, + source_agent_id=str(source_agent_id) if source_agent_id is not None else None, + message=_text(payload.get("message")), + raw=_raw_for(kind, payload), + ) + + # SubagentStart's prompt is the task the parent handed the subagent. + if kind == EventKind.SUBAGENT_START and event.prompt is None: + event.prompt = _text(_first(payload, "task", "description", "agent_prompt")) + + return event + + +def _raw_for(kind: str, payload: dict[str, Any]) -> dict[str, Any] | None: + """Keep the raw payload only for events whose fields we do not fully model. + + Tool and turn events already have every field mapped, so storing the raw + copy would double the size of the buffer for no query value. + """ + if kind in (EventKind.NOTIFICATION, EventKind.COMPACT, EventKind.SESSION_END): + return {k: v for k, v in payload.items() if k not in ("transcript_path",)} + return None + + +def handle(payload: dict[str, Any], config: Config) -> list[dict[str, Any]]: + """Record one hook payload. Returns the emitted records (for tests).""" + event = to_event(payload) + if event is None: + return [] + recorder = Recorder(config, on_error=lambda msg: log_error(config, msg)) + return recorder.record(event) + + +def main(argv: list[str] | None = None) -> int: + """Entry point for ``flowcept-claude-code`` (the hook command).""" + argv = list(sys.argv[1:] if argv is None else argv) + + def _run(config: Config) -> None: + payload = read_stdin_json() + # `--event NAME` lets one script serve every hook even on harness + # versions that omit hook_event_name from the payload. + if "--event" in argv: + idx = argv.index("--event") + if idx + 1 < len(argv): + payload["hook_event_name"] = payload.get("hook_event_name") or argv[idx + 1] + handle(payload, config) + + return run_capture(_run) + + +if __name__ == "__main__": # pragma: no cover + raise SystemExit(main()) diff --git a/src/flowcept/agents/cli_harness/cli_harness_plugin.py b/src/flowcept/agents/cli_harness/cli_harness_plugin.py new file mode 100644 index 00000000..edac0e44 --- /dev/null +++ b/src/flowcept/agents/cli_harness/cli_harness_plugin.py @@ -0,0 +1,298 @@ +"""Adapter for any harness that can run a command per lifecycle event. + +Codex CLI, Gemini CLI, Cursor, OpenCode and friends all have the same shape as +Claude Code -- a JSON event handed to a hook -- but disagree on what the fields +are called. Rather than a module per harness, the differences live in declarative +JSON *profiles* under ``profiles/``: + + { + "harness": "codex", + "events": {"tool.start": "tool_pre"}, + "fields": {"tool_name": ["tool", "name"]} + } + +``fields`` maps a normalized :class:`~flowcept.agents.harness.events.HarnessEvent` +attribute to the source keys to try, in order. Dotted keys reach into nested +objects (``"payload.tool.name"``), so a profile can flatten a nested envelope +without any code. + +Adding a harness is therefore a JSON file, and an unknown harness still works: +with no profile, the built-in default field names cover most of them, and +anything unmapped is preserved in ``custom_metadata.raw_event``. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +from flowcept.agents.harness.config import Config +from flowcept.agents.harness.events import HarnessEvent +from flowcept.agents.harness.recorder import Recorder +from flowcept.agents.harness.runtime import log_error, project_dir_from_env +from flowcept.agents.harness.vocab import EventKind + +PROFILE_DIR = Path(__file__).parent / "profiles" + +#: Tried when a profile does not name a field explicitly. Ordered by how +#: common the spelling is across harnesses. +DEFAULT_FIELDS: dict[str, list[str]] = { + "session_id": ["session_id", "sessionId", "conversation_id", "conversationId", "thread_id", "id"], + "cwd": ["cwd", "workspace", "workingDirectory", "working_dir", "project_root"], + "model": ["model", "model_id", "modelId", "model_name"], + "prompt": ["prompt", "user_prompt", "userPrompt", "input", "message", "text"], + "response": ["response", "output", "last_assistant_message", "assistant_message", "completion"], + "tool_name": ["tool_name", "toolName", "tool", "name", "function"], + "tool_use_id": ["tool_use_id", "toolUseId", "tool_call_id", "toolCallId", "call_id", "invocation_id"], + "tool_input": ["tool_input", "toolInput", "arguments", "args", "params", "input"], + "tool_response": ["tool_response", "toolResponse", "result", "output", "return_value"], + "error": ["error", "error_message", "errorMessage", "stderr", "exception"], + "agent_name": ["agent_type", "agentType", "agent_name", "subagent_type", "role"], + "agent_ref": ["agent_id", "agentId", "subagent_id", "child_session_id"], + # Flowcept's own extension key: a framework-emitted task/agent id to link + # back to. Wins over the FLOWCEPT_HARNESS_SOURCE_AGENT_ID env fallback. + "source_agent_id": ["flowcept_source_agent_id"], + "prompt_id": ["prompt_id", "promptId", "turn_id", "turnId", "message_id"], + "source": ["source", "reason", "trigger", "event_reason"], + "permission_mode": ["permission_mode", "permissionMode", "approval_mode", "mode"], + "effort": ["effort", "reasoning_effort", "reasoningEffort"], + "message": ["message", "notification", "text"], + "call_id": ["request_id", "requestId", "response_id", "generation_id"], + "usage": ["usage", "token_usage", "tokenUsage", "tokens"], +} + +#: Event names seen in the wild, mapped onto normalized kinds. Matching is +#: case-insensitive and ignores separators, so ``PreToolUse``, ``pre_tool_use`` +#: and ``tool.before`` all land in the same place. +DEFAULT_EVENTS: dict[str, str] = { + "sessionstart": EventKind.SESSION_START, + "sessionbegin": EventKind.SESSION_START, + "start": EventKind.SESSION_START, + "sessionend": EventKind.SESSION_END, + "sessionstop": EventKind.SESSION_END, + "stop": EventKind.TURN_END, + "userpromptsubmit": EventKind.PROMPT, + "prompt": EventKind.PROMPT, + "userinput": EventKind.PROMPT, + "turnstart": EventKind.PROMPT, + "turnend": EventKind.TURN_END, + "responsecomplete": EventKind.TURN_END, + "agentmessage": EventKind.TURN_END, + "pretooluse": EventKind.TOOL_PRE, + "toolstart": EventKind.TOOL_PRE, + "toolbefore": EventKind.TOOL_PRE, + "toolcall": EventKind.TOOL_PRE, + "posttooluse": EventKind.TOOL_POST, + "toolend": EventKind.TOOL_POST, + "toolafter": EventKind.TOOL_POST, + "toolresult": EventKind.TOOL_POST, + "tooldone": EventKind.TOOL_POST, + "toolerror": EventKind.TOOL_ERROR, + "toolfailed": EventKind.TOOL_ERROR, + "posttoolusefailure": EventKind.TOOL_ERROR, + "subagentstart": EventKind.SUBAGENT_START, + "agentstart": EventKind.SUBAGENT_START, + "subagentstop": EventKind.SUBAGENT_STOP, + "agentstop": EventKind.SUBAGENT_STOP, + "llmcall": EventKind.LLM_CALL, + "modelcall": EventKind.LLM_CALL, + "inference": EventKind.LLM_CALL, + "notification": EventKind.NOTIFICATION, + "precompact": EventKind.COMPACT, + "postcompact": EventKind.COMPACT, + "compact": EventKind.COMPACT, +} + + +def _normalize_event_name(name: str) -> str: + return "".join(ch for ch in name.lower() if ch.isalnum()) + + +class Profile: + """A harness's field and event naming, loaded from JSON.""" + + __slots__ = ("constants", "events", "fields", "harness") + + def __init__(self, harness: str, events=None, fields=None, constants=None): + self.harness = harness + self.events = {_normalize_event_name(k): v for k, v in (events or {}).items()} + self.fields = fields or {} + self.constants = constants or {} + + @classmethod + def load(cls, name: str | None, harness: str) -> Profile: + """Load a profile by name or path, falling back to the defaults.""" + if not name: + candidate = PROFILE_DIR / f"{harness}.json" + if not candidate.is_file(): + return cls(harness) + else: + candidate = Path(name) + if not candidate.is_file(): + candidate = PROFILE_DIR / f"{name}.json" + if not candidate.is_file(): + return cls(harness) + try: + data = json.loads(candidate.read_text(encoding="utf-8")) + except (OSError, ValueError): + return cls(harness) + return cls( + harness=data.get("harness") or harness, + events=data.get("events"), + fields=data.get("fields"), + constants=data.get("constants"), + ) + + def kind_for(self, event_name: str) -> str | None: + """Return the event kind mapped to *event_name*, or None if unmapped.""" + key = _normalize_event_name(event_name) + return self.events.get(key) or DEFAULT_EVENTS.get(key) + + def keys_for(self, field: str) -> list[str]: + """Return the payload keys to try when extracting *field*.""" + configured = self.fields.get(field) + if configured is None: + return DEFAULT_FIELDS.get(field, [field]) + if isinstance(configured, str): + return [configured] + return list(configured) + + +def _dig(payload: dict[str, Any], key: str) -> Any: + """Look up ``key``, descending through dots into nested objects. + + A numeric segment indexes a list, so ``workspace_roots.0`` reaches the + first element of a list-valued key. + """ + if "." not in key: + return payload.get(key) + current: Any = payload + for part in key.split("."): + if isinstance(current, dict): + current = current.get(part) + elif isinstance(current, (list, tuple)) and part.lstrip("-").isdigit(): + index = int(part) + current = current[index] if -len(current) <= index < len(current) else None + else: + return None + if current is None: + return None + return current + + +def _pick(payload: dict[str, Any], keys: list[str]) -> Any: + for key in keys: + value = _dig(payload, key) + if value is not None: + return value + return None + + +def _as_text(value: Any) -> str | None: + if value is None or isinstance(value, str): + return value + if isinstance(value, (int, float, bool)): + return str(value) + if isinstance(value, list): + parts = [p for p in (_as_text(v) for v in value) if p] + return "\n".join(parts) if parts else None + if isinstance(value, dict): + for key in ("text", "content", "message", "value"): + if isinstance(value.get(key), str): + return value[key] + return None + return str(value) + + +def _event_name(payload: dict[str, Any], explicit: str | None) -> str | None: + if explicit: + return explicit + for key in ("hook_event_name", "event", "event_name", "eventName", "type", "kind", "phase"): + value = payload.get(key) + if isinstance(value, str): + return value + return None + + +def to_event( + payload: dict[str, Any], + *, + harness: str, + profile: Profile, + event: str | None = None, +) -> HarnessEvent | None: + """Map an arbitrary harness payload onto a normalized event.""" + name = _event_name(payload, event) + if not name: + return None + kind = profile.kind_for(name) + if kind is None: + return None + + session_id = _as_text(_pick(payload, profile.keys_for("session_id"))) + if not session_id: + return None + + cwd = _as_text(_pick(payload, profile.keys_for("cwd"))) + usage = _pick(payload, profile.keys_for("usage")) + + normalized = HarnessEvent( + kind=kind, + harness=profile.harness or harness, + session_id=session_id, + cwd=cwd, + project_dir=project_dir_from_env(cwd), + model=_as_text(_pick(payload, profile.keys_for("model"))), + permission_mode=_as_text(_pick(payload, profile.keys_for("permission_mode"))), + effort=_as_text(_pick(payload, profile.keys_for("effort"))), + source=_as_text(_pick(payload, profile.keys_for("source"))), + prompt_id=_as_text(_pick(payload, profile.keys_for("prompt_id"))), + error=_as_text(_pick(payload, profile.keys_for("error"))), + call_id=_as_text(_pick(payload, profile.keys_for("call_id"))), + usage=usage if isinstance(usage, dict) else None, + agent_name=_as_text(_pick(payload, profile.keys_for("agent_name"))), + agent_ref=_as_text(_pick(payload, profile.keys_for("agent_ref"))), + source_agent_id=_as_text(_pick(payload, profile.keys_for("source_agent_id"))), + message=_as_text(_pick(payload, profile.keys_for("message"))), + # Only lifecycle events keep their raw payload (the recorder discards it + # otherwise). Tool and turn events are high-volume and would double the + # buffer; their fields are the ones profiles exist to map anyway. + raw=payload if kind in (EventKind.NOTIFICATION, EventKind.COMPACT, EventKind.SESSION_END) else None, + ) + + # Prompt and tool fields share source keys across harnesses ("input" is a + # prompt on a turn event and arguments on a tool event), so they are only + # read for the events where they mean what we want. + if kind in (EventKind.PROMPT, EventKind.TURN_END, EventKind.LLM_CALL, EventKind.SUBAGENT_START): + normalized.prompt = _as_text(_pick(payload, profile.keys_for("prompt"))) + if kind in (EventKind.TURN_END, EventKind.LLM_CALL, EventKind.SUBAGENT_STOP): + normalized.response = _as_text(_pick(payload, profile.keys_for("response"))) + if kind in (EventKind.TOOL_PRE, EventKind.TOOL_POST, EventKind.TOOL_ERROR): + normalized.tool_name = _as_text(_pick(payload, profile.keys_for("tool_name"))) + normalized.tool_use_id = _as_text(_pick(payload, profile.keys_for("tool_use_id"))) + normalized.tool_input = _pick(payload, profile.keys_for("tool_input")) + normalized.tool_response = _pick(payload, profile.keys_for("tool_response")) + + for key, value in profile.constants.items(): + if hasattr(normalized, key): + setattr(normalized, key, value) + + return normalized + + +def handle( + payload: dict[str, Any], + config: Config, + *, + harness: str = "generic", + profile: str | None = None, + event: str | None = None, +) -> list[dict[str, Any]]: + """Record one payload from an arbitrary harness.""" + loaded = Profile.load(profile, harness) + normalized = to_event(payload, harness=harness, profile=loaded, event=event) + if normalized is None: + return [] + recorder = Recorder(config, on_error=lambda msg: log_error(config, msg)) + return recorder.record(normalized) diff --git a/src/flowcept/agents/cli_harness/profiles/codex.json b/src/flowcept/agents/cli_harness/profiles/codex.json new file mode 100644 index 00000000..9b642c35 --- /dev/null +++ b/src/flowcept/agents/cli_harness/profiles/codex.json @@ -0,0 +1,32 @@ +{ + "harness": "codex", + "_comment": "OpenAI Codex CLI. Notification hooks deliver a typed JSON envelope; the turn/tool names below cover the documented types. Unmapped names fall through to the built-in defaults.", + "events": { + "session-start": "session_start", + "session-end": "session_end", + "session-configured": "session_start", + "user-message": "prompt", + "agent-turn-complete": "turn_end", + "agent-message": "turn_end", + "exec-command-begin": "tool_pre", + "exec-command-end": "tool_post", + "patch-apply-begin": "tool_pre", + "patch-apply-end": "tool_post", + "mcp-tool-call-begin": "tool_pre", + "mcp-tool-call-end": "tool_post", + "error": "tool_error" + }, + "fields": { + "session_id": ["session_id", "conversation_id", "id"], + "cwd": ["cwd", "workdir", "turn_context.cwd"], + "model": ["model", "turn_context.model"], + "prompt": ["input_messages", "last_user_message", "message"], + "response": ["last_agent_message", "agent_message", "message"], + "tool_name": ["invocation.tool", "tool", "command", "call.name"], + "tool_use_id": ["call_id", "invocation.call_id", "id"], + "tool_input": ["invocation.arguments", "command", "changes", "arguments"], + "tool_response": ["output", "stdout", "result", "aggregated_output"], + "error": ["error", "stderr", "message"], + "permission_mode": ["turn_context.approval_policy", "approval_policy", "sandbox_policy"] + } +} diff --git a/src/flowcept/agents/cli_harness/profiles/cursor.json b/src/flowcept/agents/cli_harness/profiles/cursor.json new file mode 100644 index 00000000..2ec22375 --- /dev/null +++ b/src/flowcept/agents/cli_harness/profiles/cursor.json @@ -0,0 +1,32 @@ +{ + "harness": "cursor", + "_comment": "Cursor agent hooks. Cursor identifies a conversation with `conversation_id` and the workspace with `workspace_roots`; tool events carry `tool_name` plus a free-form `tool_input`.", + "events": { + "beforeSubmitPrompt": "prompt", + "afterPromptSubmit": "prompt", + "beforeShellExecution": "tool_pre", + "beforeMCPExecution": "tool_pre", + "beforeReadFile": "tool_pre", + "beforeEdit": "tool_pre", + "afterFileEdit": "tool_post", + "afterShellExecution": "tool_post", + "afterMCPExecution": "tool_post", + "stop": "turn_end", + "afterAgentTurn": "turn_end", + "start": "session_start", + "sessionEnd": "session_end" + }, + "fields": { + "session_id": ["conversation_id", "conversationId", "session_id", "thread_id"], + "cwd": ["workspace_roots.0", "workspace_root", "cwd"], + "model": ["model", "model_name"], + "prompt": ["prompt", "text", "attachments.prompt"], + "response": ["text", "response", "assistant_message"], + "tool_name": ["tool_name", "command", "tool", "server_name"], + "tool_use_id": ["generation_id", "tool_call_id", "call_id"], + "tool_input": ["tool_input", "command", "edits", "args"], + "tool_response": ["tool_response", "output", "result"], + "error": ["error", "stderr"], + "permission_mode": ["permission", "mode"] + } +} diff --git a/src/flowcept/agents/cli_harness/profiles/gemini.json b/src/flowcept/agents/cli_harness/profiles/gemini.json new file mode 100644 index 00000000..cb6be0af --- /dev/null +++ b/src/flowcept/agents/cli_harness/profiles/gemini.json @@ -0,0 +1,31 @@ +{ + "harness": "gemini_cli", + "_comment": "Google Gemini CLI. Its extension/hook events are camelCase and nest tool detail under a `toolCall` object.", + "events": { + "SessionStart": "session_start", + "SessionEnd": "session_end", + "UserPromptSubmit": "prompt", + "BeforeToolCall": "tool_pre", + "PreToolUse": "tool_pre", + "AfterToolCall": "tool_post", + "PostToolUse": "tool_post", + "ToolCallError": "tool_error", + "ModelResponse": "turn_end", + "AgentFinish": "turn_end", + "Notification": "notification", + "Compress": "compact" + }, + "fields": { + "session_id": ["sessionId", "session_id", "conversationId"], + "cwd": ["cwd", "workspaceDir", "targetDir"], + "model": ["model", "modelName", "config.model"], + "prompt": ["prompt", "userPrompt", "text"], + "response": ["responseText", "modelResponse", "text"], + "tool_name": ["toolCall.name", "toolName", "name"], + "tool_use_id": ["toolCall.callId", "callId", "toolCallId"], + "tool_input": ["toolCall.args", "args", "toolArgs"], + "tool_response": ["toolCall.response", "response", "resultDisplay"], + "error": ["error", "errorMessage", "toolCall.error"], + "permission_mode": ["approvalMode", "config.approvalMode"] + } +} diff --git a/src/flowcept/agents/cli_harness/profiles/opencode.json b/src/flowcept/agents/cli_harness/profiles/opencode.json new file mode 100644 index 00000000..c3dec549 --- /dev/null +++ b/src/flowcept/agents/cli_harness/profiles/opencode.json @@ -0,0 +1,28 @@ +{ + "harness": "opencode", + "_comment": "OpenCode. Its plugin API emits dotted event names and wraps detail in a nested object, so most lookups are dotted paths.", + "events": { + "session.created": "session_start", + "session.idle": "turn_end", + "session.deleted": "session_end", + "session.error": "turn_end", + "message.updated": "prompt", + "tool.execute.before": "tool_pre", + "tool.execute.after": "tool_post", + "permission.updated": "notification", + "session.compacted": "compact" + }, + "fields": { + "session_id": ["sessionID", "properties.sessionID", "properties.info.id", "session_id"], + "cwd": ["directory", "properties.directory", "worktree"], + "model": ["properties.model", "model", "properties.info.modelID"], + "prompt": ["properties.text", "properties.info.text", "text"], + "response": ["properties.text", "properties.info.summary", "text"], + "tool_name": ["tool", "properties.tool", "name"], + "tool_use_id": ["callID", "properties.callID", "sessionID"], + "tool_input": ["args", "properties.args", "input"], + "tool_response": ["output", "properties.output", "result"], + "error": ["error", "properties.error", "properties.info.error"], + "agent_name": ["agent", "properties.agent", "properties.info.agent"] + } +} diff --git a/src/flowcept/agents/crewai/crewai_plugin.py b/src/flowcept/agents/crewai/crewai_plugin.py new file mode 100644 index 00000000..f8a9379e --- /dev/null +++ b/src/flowcept/agents/crewai/crewai_plugin.py @@ -0,0 +1,1396 @@ +# academy_coscientist/plugins/flowcept_crewai_plugin.py +""" +FlowCept provenance plugin for CrewAI workflows. + +Captures the full provenance graph for each crew kickoff using two complementary +CrewAI extension points: + + 1. Event bus (BaseEventListener) — crew kickoff, task, and agent lifecycle events. + 2. LLM + Tool hooks (register_before/after_llm_call_hook, + register_before/after_tool_call_hook) — richer LLM and tool provenance + including full messages, response text, agent.role/goal/backstory, + task.description/expected_output, iteration count, and typed tool inputs. + +The event bus is used for lifecycle events (crew, task, agent). +The hooks replace the thin LLMCallStarted/Completed events with richer records. + +Provenance hierarchy produced: + WorkflowObject (one per crew kickoff) + └─ TaskObject subtype=crewai_task activity_id= + └─ TaskObject subtype=crewai_agent activity_id= + └─ TaskObject subtype=llm_call activity_id= + └─ TaskObject subtype=tool_call activity_id= + +Key FlowCept fields: + task_id — uuid per event (our own, not CrewAI's event_id) + workflow_id — global shared workflow id (setdefault pattern) + campaign_id — from Flowcept.campaign_id + parent_task_id — links agent → task, llm/tool → agent + group_id — all tasks within one crew kickoff share a group_id + activity_id — task name | agent role | model name | tool name + subtype — crewai_crew | crewai_task | crewai_agent | llm_call | tool_call + used / generated — inputs and outputs + status — FINISHED | ERROR + +Usage (two lines to wire up): + + plugin = FlowceptCrewAIPlugin(config={"workflow_name": "my-crew"}) + plugin.start() + crew.kickoff() + plugin.stop() + +Or shared with an Academy plugin: + + academy_plugin = FlowceptAcademyPlugin(config={...}).start() + crewai_plugin = FlowceptCrewAIPlugin.from_academy_plugin(academy_plugin) + crew.kickoff() + academy_plugin.stop() # flushes the shared buffer +""" + +from __future__ import annotations + +import os +import time +import uuid +import threading +import logging +from contextlib import contextmanager +from typing import Any + +_log = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Provenance overhead timer (identical to other plugins) +# --------------------------------------------------------------------------- + + +class _ProvenanceStats: + __slots__ = ("_lock", "_counts", "_totals", "_mins", "_maxs", "_raw") + + def __init__(self) -> None: + self._lock: threading.Lock = threading.Lock() + self._counts: dict[str, int] = {} + self._totals: dict[str, float] = {} + self._mins: dict[str, float] = {} + self._maxs: dict[str, float] = {} + self._raw: list[tuple[str, str, float]] = [] + + def record(self, category: str, elapsed: float) -> None: + ts = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) + with self._lock: + if category not in self._counts: + self._counts[category] = 0 + self._totals[category] = 0.0 + self._mins[category] = float("inf") + self._maxs[category] = 0.0 + self._counts[category] += 1 + self._totals[category] += elapsed + if elapsed < self._mins[category]: + self._mins[category] = elapsed + if elapsed > self._maxs[category]: + self._maxs[category] = elapsed + self._raw.append((ts, category, elapsed)) + + def summary(self) -> str: + col = 22 + header = f"{'Category':<{col}} {'N':>7} {'Total(ms)':>11} {'Mean(µs)':>9} {'Min(µs)':>8} {'Max(µs)':>8}" + sep = "-" * len(header) + rows = [header, sep] + with self._lock: + for cat in sorted(self._counts): + n = self._counts[cat] + total = self._totals[cat] + mean = (total / n) if n else 0.0 + mn = self._mins.get(cat, 0.0) + mx = self._maxs.get(cat, 0.0) + rows.append( + f"{cat:<{col}} {n:>7} {total * 1e3:>11.3f} {mean * 1e6:>9.1f} {mn * 1e6:>8.1f} {mx * 1e6:>8.1f}" + ) + return "\n".join(rows) + + def to_csv(self, path: str, workflow_id: str | None = None) -> None: + import csv + + write_header = not os.path.exists(path) + with self._lock: + raw_snapshot = list(self._raw) + wf = workflow_id or "" + rows = [ + { + "timestamp_utc": ts, + "workflow_id": wf, + "category": cat, + "elapsed_us": round(elapsed * 1e6, 3), + } + for ts, cat, elapsed in raw_snapshot + ] + fieldnames = ["timestamp_utc", "workflow_id", "category", "elapsed_us"] + with open(path, "a", newline="", encoding="utf-8") as fh: + writer = csv.DictWriter(fh, fieldnames=fieldnames) + if write_header: + writer.writeheader() + writer.writerows(rows) + + +# --------------------------------------------------------------------------- +# CrewAI event listener — uses the native event bus +# --------------------------------------------------------------------------- + + +class _FlowceptCrewAIListener: + """ + Listens to CrewAI's event bus and records provenance to a shared interceptor. + + Registered event pairs (start→end → ONE complete record each): + CrewKickoffStarted/Completed → crewai_crew TaskObject + WorkflowObject + TaskStarted/Completed → crewai_task TaskObject + AgentExecutionStarted/Completed → crewai_agent TaskObject + LLMCallStarted/Completed → llm_call TaskObject + ToolUsageStarted/Finished → tool_call TaskObject + """ + + def __init__(self, interceptor: Any, stats: _ProvenanceStats | None) -> None: + self._interceptor = interceptor + self._stats = stats + # Buffer: crew_kickoff event_id → {group_id, started_at, task_id, used, ...} + self._crew_starts: dict[str, dict] = {} + # crewai task event_id → {task_id, started_at, ...} + self._task_starts: dict[str, dict] = {} + # agent event_id → {task_id, started_at, ...} + self._agent_starts: dict[str, dict] = {} + # llm call_id → {task_id, started_at, ...} + self._llm_starts: dict[str, dict] = {} + # tool started_event_id → {task_id, started_at, ...} + self._tool_starts: dict[str, dict] = {} + # event_id → group_id for hierarchy linking + self._crew_group: dict[str, str] = {} + # task event_id → FlowCept task_id (so agent can reference it) + self._task_fc_id: dict[str, str] = {} + # agent event_id → FlowCept task_id (so llm/tool can reference it) + self._agent_fc_id: dict[str, str] = {} + + def _record(self, category: str, elapsed: float) -> None: + if self._stats is not None: + self._stats.record(category, elapsed) + + # ── helpers ───────────────────────────────────────────────────────────── + + def _safe_clip(self, obj: Any, depth: int = 0) -> Any: + if depth > 6: + return str(obj) + if obj is None or isinstance(obj, (bool, int, float, str)): + return obj + if isinstance(obj, dict): + return {str(k): self._safe_clip(v, depth + 1) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return [self._safe_clip(v, depth + 1) for v in obj] + # For CrewAI objects: try model_dump, then repr + try: + if hasattr(obj, "model_dump"): + return self._safe_clip(obj.model_dump(), depth + 1) + except Exception: + pass + return repr(obj) + + def _intercept(self, task: dict) -> None: + self._interceptor.intercept_task(task) + + # ── Crew Kickoff ───────────────────────────────────────────────────────── + + def on_crew_kickoff_started(self, source: Any, event: Any) -> None: + t0 = time.perf_counter() + group_id = str(uuid.uuid4()) + task_id = str(uuid.uuid4()) + crew_name = getattr(event, "crew_name", None) or "crew" + inputs = self._safe_clip(getattr(event, "inputs", None) or {}) + event_id = getattr(event, "event_id", str(uuid.uuid4())) + + # Emit a sub-WorkflowObject for this kickoff so queries can navigate + # the crew hierarchy + self._interceptor.send_graph_workflow(crew_name, group_id) + self._crew_group[event_id] = group_id + + self._crew_starts[event_id] = { + "task_id": task_id, + "subtype": "crewai_crew", + "activity_id": crew_name, + "group_id": group_id, + "started_at": time.time(), + "used": {"inputs": inputs, "crew_name": crew_name}, + "custom_metadata": {"crew_name": crew_name, "framework": "crewai"}, + } + self._record("crew_kickoff_started", time.perf_counter() - t0) + + def on_crew_kickoff_completed(self, source: Any, event: Any) -> None: + t0 = time.perf_counter() + event_id = getattr(event, "event_id", None) + # CrewAI's completed event has started_event_id pointing back to start + started_id = getattr(event, "started_event_id", None) or event_id + skeleton = self._crew_starts.pop(started_id, None) or self._crew_starts.pop(event_id, {}) + + output = self._safe_clip(getattr(event, "output", None)) + total_tokens = getattr(event, "total_tokens", None) + + skeleton.update( + { + "ended_at": time.time(), + "status": "FINISHED", + "generated": {"output": output, "total_tokens": total_tokens}, + } + ) + self._intercept(skeleton) + # Clean up group mapping + self._crew_group.pop(started_id, None) + self._crew_group.pop(event_id, None) + self._record("crew_kickoff_completed", time.perf_counter() - t0) + + def on_crew_kickoff_failed(self, source: Any, event: Any) -> None: + t0 = time.perf_counter() + event_id = getattr(event, "event_id", None) + started_id = getattr(event, "started_event_id", None) or event_id + skeleton = self._crew_starts.pop(started_id, {}) + skeleton.update( + { + "ended_at": time.time(), + "status": "ERROR", + "stderr": str(getattr(event, "error", "unknown error")), + } + ) + self._intercept(skeleton) + self._record("crew_kickoff_failed", time.perf_counter() - t0) + + # ── Task ───────────────────────────────────────────────────────────────── + + def _find_group_id(self, event: Any) -> str | None: + """Walk parent event chain to find the enclosing crew's group_id.""" + # CrewAI propagates task_name / agent_id through all events — + # the group_id can be found via the triggered_by_event_id chain. + # Simplest: scan _crew_group values (there's usually one active run). + if self._crew_group: + return next(iter(self._crew_group.values())) + return None + + def on_task_started(self, source: Any, event: Any) -> None: + t0 = time.perf_counter() + event_id = getattr(event, "event_id", str(uuid.uuid4())) + task_id = str(uuid.uuid4()) + task_name = getattr(event, "task_name", None) or "task" + agent_role = getattr(event, "agent_role", None) or "unknown" + group_id = self._find_group_id(event) + context = self._safe_clip(getattr(event, "context", None) or {}) + task_obj = self._safe_clip(getattr(event, "task", None)) + + self._task_starts[event_id] = { + "task_id": task_id, + "subtype": "crewai_task", + "activity_id": task_name, + "group_id": group_id, + "started_at": time.time(), + "used": {"context": context, "task": task_obj}, + "custom_metadata": { + "task_name": task_name, + "agent_role": agent_role, + "framework": "crewai", + }, + } + self._task_fc_id[event_id] = task_id + self._record("task_started", time.perf_counter() - t0) + + def on_task_completed(self, source: Any, event: Any) -> None: + t0 = time.perf_counter() + event_id = getattr(event, "event_id", None) + started_id = getattr(event, "started_event_id", None) or event_id + skeleton = self._task_starts.pop(started_id, None) or self._task_starts.pop(event_id, {}) + + output = self._safe_clip(getattr(event, "output", None)) + skeleton.update( + { + "ended_at": time.time(), + "status": "FINISHED", + "generated": {"output": output}, + } + ) + self._intercept(skeleton) + self._task_fc_id.pop(started_id, None) + self._task_fc_id.pop(event_id, None) + self._record("task_completed", time.perf_counter() - t0) + + def on_task_failed(self, source: Any, event: Any) -> None: + t0 = time.perf_counter() + event_id = getattr(event, "event_id", None) + started_id = getattr(event, "started_event_id", None) or event_id + skeleton = self._task_starts.pop(started_id, {}) + skeleton.update( + { + "ended_at": time.time(), + "status": "ERROR", + "stderr": str(getattr(event, "error", "unknown")), + } + ) + self._intercept(skeleton) + self._record("task_failed", time.perf_counter() - t0) + + # ── Agent Execution ─────────────────────────────────────────────────────── + + def on_agent_execution_started(self, source: Any, event: Any) -> None: + t0 = time.perf_counter() + event_id = getattr(event, "event_id", str(uuid.uuid4())) + task_id = str(uuid.uuid4()) + agent = getattr(event, "agent", None) + agent_role = getattr(event, "agent_role", None) or (getattr(agent, "role", None) if agent else None) or "agent" + group_id = self._find_group_id(event) + task_prompt = self._safe_clip(getattr(event, "task_prompt", None) or "") + tools = self._safe_clip([getattr(t, "name", str(t)) for t in (getattr(event, "tools", None) or [])]) + # Link to enclosing task's FlowCept task_id + task_evt_id = getattr(event, "started_event_id", None) or getattr(event, "task_id", None) + parent_task_id = self._task_fc_id.get(task_evt_id) if task_evt_id else None + + skeleton: dict = { + "task_id": task_id, + "subtype": "crewai_agent", + "activity_id": agent_role, + "group_id": group_id, + "started_at": time.time(), + "used": {"task_prompt": task_prompt, "tools": tools}, + "custom_metadata": { + "agent_role": agent_role, + "framework": "crewai", + }, + } + if parent_task_id: + skeleton["parent_task_id"] = parent_task_id + self._agent_starts[event_id] = skeleton + self._agent_fc_id[event_id] = task_id + self._record("agent_started", time.perf_counter() - t0) + + def on_agent_execution_completed(self, source: Any, event: Any) -> None: + t0 = time.perf_counter() + event_id = getattr(event, "event_id", None) + started_id = getattr(event, "started_event_id", None) or event_id + skeleton = self._agent_starts.pop(started_id, None) or self._agent_starts.pop(event_id, {}) + + output = self._safe_clip(getattr(event, "output", None)) + skeleton.update( + { + "ended_at": time.time(), + "status": "FINISHED", + "generated": {"output": output}, + } + ) + self._intercept(skeleton) + self._agent_fc_id.pop(started_id, None) + self._agent_fc_id.pop(event_id, None) + self._record("agent_completed", time.perf_counter() - t0) + + def on_agent_execution_error(self, source: Any, event: Any) -> None: + t0 = time.perf_counter() + event_id = getattr(event, "event_id", None) + started_id = getattr(event, "started_event_id", None) or event_id + skeleton = self._agent_starts.pop(started_id, {}) + skeleton.update( + { + "ended_at": time.time(), + "status": "ERROR", + "stderr": str(getattr(event, "error", "unknown")), + } + ) + self._intercept(skeleton) + self._record("agent_error", time.perf_counter() - t0) + + # ── LLM + Tool calls are captured via hooks (see _FlowceptCrewAIHooks) ──── + # The event bus LLM/tool events are deliberately NOT registered here; + # hooks provide richer data (full messages, agent/task objects, response text). + + +def _build_listener_class(interceptor: Any, stats: _ProvenanceStats | None): + """ + Lazily build a concrete BaseEventListener subclass. + + Ensures crewai is only imported when the plugin is actually used. + Registers only lifecycle events (crew, task, agent). + LLM and tool calls are captured via hooks (richer data). + """ + from crewai.events.base_event_listener import BaseEventListener + from crewai.events.event_types import ( + CrewKickoffStartedEvent, + CrewKickoffCompletedEvent, + CrewKickoffFailedEvent, + TaskStartedEvent, + TaskCompletedEvent, + TaskFailedEvent, + AgentExecutionStartedEvent, + AgentExecutionCompletedEvent, + AgentExecutionErrorEvent, + ) + + _listener = _FlowceptCrewAIListener(interceptor, stats) + + class _ConcreteListener(BaseEventListener): + def setup_listeners(self, bus) -> None: + bus.on(CrewKickoffStartedEvent)(_listener.on_crew_kickoff_started) + bus.on(CrewKickoffCompletedEvent)(_listener.on_crew_kickoff_completed) + bus.on(CrewKickoffFailedEvent)(_listener.on_crew_kickoff_failed) + bus.on(TaskStartedEvent)(_listener.on_task_started) + bus.on(TaskCompletedEvent)(_listener.on_task_completed) + bus.on(TaskFailedEvent)(_listener.on_task_failed) + bus.on(AgentExecutionStartedEvent)(_listener.on_agent_execution_started) + bus.on(AgentExecutionCompletedEvent)(_listener.on_agent_execution_completed) + bus.on(AgentExecutionErrorEvent)(_listener.on_agent_execution_error) + + return _listener, _ConcreteListener() + + +# --------------------------------------------------------------------------- +# Hook-based LLM + Tool provenance (richer than event bus events) +# --------------------------------------------------------------------------- + + +class _FlowceptCrewAIHooks: + """ + Registers CrewAI's before/after LLM and tool hooks to capture rich provenance. + + LLM hooks give: + - Full message list (context.messages) — not just serialised dicts + - Actual response text (context.response) + - agent.role, agent.goal, agent.backstory + - task.description, task.expected_output + - iterations (current ReAct loop count) + + Tool hooks give: + - Typed tool_input dict (not a string) + - tool_result string + - agent, task, crew references + """ + + def __init__( + self, + interceptor: Any, + stats: _ProvenanceStats | None, + listener: _FlowceptCrewAIListener, + ) -> None: + self._interceptor = interceptor + self._stats = stats + self._listener = listener # shares group_id / agent_fc_id state + # Keyed by id(context.executor): before hook stores skeleton, + # after hook pops and emits. CrewAI creates new context objects for + # each hook invocation so the context itself cannot be used as scratchpad. + self._pending_llm: dict[int, dict] = {} + self._pending_tool: dict[int, dict] = {} + + def _record(self, category: str, elapsed: float) -> None: + if self._stats is not None: + self._stats.record(category, elapsed) + + def _safe_clip(self, obj: Any, depth: int = 0) -> Any: + return self._listener._safe_clip(obj, depth) + + # ── Before-LLM hook ────────────────────────────────────────────────────── + + def before_llm_call(self, context: Any) -> None: + """Buffers the LLM call start skeleton; matched by after_llm_call.""" + t0 = time.perf_counter() + task_id = str(uuid.uuid4()) + group_id = self._listener._find_group_id(None) + + # Rich context from the hook + agent = getattr(context, "agent", None) + task_obj = getattr(context, "task", None) + llm = getattr(context, "llm", None) + messages = getattr(context, "messages", []) + iterations = getattr(context, "iterations", 0) + + model = getattr(llm, "model", None) or getattr(llm, "model_name", None) or "unknown" + agent_role = getattr(agent, "role", None) or "unknown" + + # Retrieve enclosing agent's FlowCept task_id for parent linkage + # The agent_fc_id map is keyed by CrewAI event_id; best-effort lookup + parent_task_id = next(iter(self._listener._agent_fc_id.values()), None) if self._listener._agent_fc_id else None + + serialized_messages = self._safe_clip( + [{"role": getattr(m, "role", "user"), "content": getattr(m, "content", str(m))} for m in messages] + if messages and hasattr(messages[0], "role") + else messages + ) + + used: dict = { + "messages": serialized_messages, + "model": model, + "iterations": iterations, + "agent_role": agent_role, + "agent_goal": self._safe_clip(getattr(agent, "goal", None)), + "task_description": self._safe_clip(getattr(task_obj, "description", None)), + } + + skeleton: dict = { + "task_id": task_id, + "subtype": "llm_call", + "activity_id": model, + "group_id": group_id, + "started_at": time.time(), + "used": used, + "custom_metadata": { + "model": model, + "agent_role": agent_role, + "framework": "crewai", + "source": "llm_hook", + }, + } + if parent_task_id: + skeleton["parent_task_id"] = parent_task_id + + # CrewAI creates NEW context objects for before vs after hooks so we + # cannot store state on the context. Key by executor id instead — + # each executor handles one LLM call at a time (sequential). + key = id(getattr(context, "executor", context)) + self._pending_llm[key] = skeleton + self._record("llm_hook_before", time.perf_counter() - t0) + + # ── After-LLM hook ─────────────────────────────────────────────────────── + + def after_llm_call(self, context: Any) -> None: + """Completes the skeleton buffered by before_llm_call and emits it.""" + t0 = time.perf_counter() + key = id(getattr(context, "executor", context)) + skeleton: dict = self._pending_llm.pop(key, {}) + response = getattr(context, "response", None) or "" + skeleton.update( + { + "ended_at": time.time(), + "status": "FINISHED", + "generated": { + "response": response[:4000] if isinstance(response, str) else self._safe_clip(response), + "model": skeleton.get("activity_id", "unknown"), + }, + } + ) + self._interceptor.intercept_task(skeleton) + self._record("llm_hook_after", time.perf_counter() - t0) + + # ── Before-tool hook ───────────────────────────────────────────────────── + + def before_tool_call(self, context: Any) -> None: + """Buffers the tool call start skeleton.""" + t0 = time.perf_counter() + task_id = str(uuid.uuid4()) + tool_name = getattr(context, "tool_name", None) or "tool" + group_id = self._listener._find_group_id(None) + + agent = getattr(context, "agent", None) + task_obj = getattr(context, "task", None) + agent_role = getattr(agent, "role", None) or "unknown" + + parent_task_id = next(iter(self._listener._agent_fc_id.values()), None) if self._listener._agent_fc_id else None + + skeleton: dict = { + "task_id": task_id, + "subtype": "tool_call", + "activity_id": tool_name, + "group_id": group_id, + "started_at": time.time(), + "used": { + "input": self._safe_clip(getattr(context, "tool_input", {})), + "agent_role": agent_role, + "task_description": self._safe_clip(getattr(task_obj, "description", None)), + }, + "custom_metadata": { + "tool_name": tool_name, + "agent_role": agent_role, + "framework": "crewai", + "source": "tool_hook", + }, + } + if parent_task_id: + skeleton["parent_task_id"] = parent_task_id + + key = id(getattr(context, "executor", context)) + self._pending_tool[key] = skeleton + self._record("tool_hook_before", time.perf_counter() - t0) + + # ── After-tool hook ────────────────────────────────────────────────────── + + def after_tool_call(self, context: Any) -> None: + """Completes the skeleton buffered by before_tool_call and emits it.""" + t0 = time.perf_counter() + key = id(getattr(context, "executor", context)) + skeleton: dict = self._pending_tool.pop(key, {}) + result = getattr(context, "tool_result", None) or "" + skeleton.update( + { + "ended_at": time.time(), + "status": "FINISHED", + "generated": {"output": result[:2000] if isinstance(result, str) else self._safe_clip(result)}, + } + ) + self._interceptor.intercept_task(skeleton) + self._record("tool_hook_after", time.perf_counter() - t0) + + +# --------------------------------------------------------------------------- +# Shared interceptor wrapper (re-uses AcademyInterceptor when available) +# --------------------------------------------------------------------------- + + +class _CrewAIInterceptor: + """Standalone interceptor for CrewAI provenance (Dask-style).""" + + def __init__(self) -> None: + self._interceptor = None + self._workflow_id: str | None = None + self._campaign_id: str | None = None + + def start(self, workflow_name: str, campaign_id: str | None = None) -> None: + from flowcept.flowceptor.adapters.base_interceptor import BaseInterceptor + from flowcept.commons.flowcept_dataclasses.workflow_object import WorkflowObject + + self._workflow_id = str(uuid.uuid4()) + self._campaign_id = campaign_id or str(uuid.uuid4()) + + self._interceptor = BaseInterceptor(kind="crewai") + self._interceptor.start( + bundle_exec_id=self._workflow_id, + check_safe_stops=False, + ) + + wf = WorkflowObject() + wf.workflow_id = self._workflow_id + wf.campaign_id = self._campaign_id + wf.name = workflow_name + self._interceptor.send_workflow_message(wf) + + def stop(self) -> None: + if self._interceptor is None: + return + try: + self._interceptor.stop(check_safe_stops=False) + except Exception as e: + _log.warning("Interceptor stop error: %r", e) + self._interceptor = None + + def send_graph_workflow(self, name: str, group_id: str) -> str: + if self._interceptor is None: + return str(uuid.uuid4()) + from flowcept.commons.flowcept_dataclasses.workflow_object import WorkflowObject + + wf = WorkflowObject() + wf.workflow_id = str(uuid.uuid4()) + wf.name = name + wf.campaign_id = self._campaign_id + wf.parent_workflow_id = self._workflow_id + wf.custom_metadata = {"group_id": group_id, "framework": "crewai"} + self._interceptor.send_workflow_message(wf) + return wf.workflow_id + + def intercept_task(self, task_dict: dict) -> None: + if self._interceptor is None: + return + from flowcept.commons.flowcept_dataclasses.task_object import TaskObject + from flowcept.commons.vocabulary import Status + + task_dict.setdefault("task_id", str(uuid.uuid4())) + task_dict.setdefault("workflow_id", self._workflow_id) + task_dict.setdefault("campaign_id", self._campaign_id) + + raw = task_dict.get("status", "FINISHED") + if isinstance(raw, str): + try: + task_dict["status"] = Status[raw].value + except KeyError: + task_dict["status"] = Status.FINISHED.value + + TaskObject.enrich_task_dict(task_dict) + self._interceptor.intercept(task_dict) + + +# --------------------------------------------------------------------------- +# Module-level active interceptor — set by start() / from_academy_plugin() +# --------------------------------------------------------------------------- + +_ACTIVE_INTERCEPTOR = None +_PROV_STATS: _ProvenanceStats | None = None + + +@contextmanager +def _timed(category: str): + t0 = time.perf_counter() + try: + yield + finally: + if _PROV_STATS is not None: + _PROV_STATS.record(category, time.perf_counter() - t0) + + +def record_llm_call(payload: dict) -> None: + """ + Public API to record an LLM call into the active FlowCept provenance graph. + + Converts the payload into a TaskObject (subtype=llm_call) and routes it + through the active interceptor. No-ops if the plugin has not been started. + + Minimum payload keys: + type : "chat_completion" + model : str + text : str + usage : dict with prompt_tokens / completion_tokens / total_tokens + """ + interceptor = _ACTIVE_INTERCEPTOR + if interceptor is None: + return + import uuid as _uuid + + with _timed("record_llm_call"): + elapsed = payload.get("elapsed_s", 0.0) + now = time.time() + model = payload.get("model_used") or payload.get("model", "unknown") + + used: dict = {} + for k in ( + "model", + "model_used", + "messages", + "user_prompt", + "system_prompt", + "temperature", + "top_p", + "max_tokens", + "reasoning_effort", + "tools_provided", + "tool_choice", + "stop_sequences", + "top_k", + "thinking_budget_tokens", + ): + if k in payload: + used[k] = payload[k] + if payload.get("temperature_suppressed"): + used["temperature_suppressed"] = True + + generated: dict = {} + for k in ( + "text", + "finish_reason", + "stop_reason", + "stop_sequence", + "tool_calls", + "tool_uses", + "thinking_text", + "system_fingerprint", + "response_id", + "usage", + "elapsed_s", + ): + if k in payload: + generated[k] = payload[k] + if "error" in payload: + generated["error"] = str(payload["error"]) + + task: dict = { + "task_id": str(_uuid.uuid4()), + "subtype": "llm_call", + "activity_id": model, + "started_at": now - elapsed, + "ended_at": now, + "status": "ERROR" if "error" in payload else "FINISHED", + "used": used, + "generated": generated, + "custom_metadata": { + "model": model, + "framework": payload.get("context", {}).get("framework", ""), + "context": payload.get("context", {}), + }, + } + interceptor.intercept_task(task) + + +def openai_chat( + prompt: str, + model: str = "gpt-4o-mini", + system: str = "You are a helpful assistant.", + temperature: float | None = 0.3, + top_p: float | None = None, + max_tokens: int | None = None, + n: int = 1, + stop: list[str] | str | None = None, + frequency_penalty: float = 0.0, + presence_penalty: float = 0.0, + seed: int | None = None, + reasoning_effort: str | None = None, + response_format: dict | None = None, + tools: list | None = None, + tool_choice: str | dict | None = None, + user: str | None = None, + context: dict | None = None, +) -> str: + """ + Make an OpenAI chat completion call and record it for FlowCept provenance. + + Captures all request parameters and response fields as a child TaskObject. + No-ops gracefully if OPENAI_API_KEY is not set or the plugin is not started. + """ + import openai as _openai + import os as _os + import time as _time + + client = _openai.OpenAI(api_key=_os.environ.get("OPENAI_API_KEY")) + messages = [ + {"role": "system", "content": system}, + {"role": "user", "content": prompt}, + ] + req: dict = {"model": model, "messages": messages, "n": n} + if temperature is not None: + req["temperature"] = temperature + if top_p is not None: + req["top_p"] = top_p + if max_tokens is not None: + req["max_completion_tokens"] = max_tokens + if stop is not None: + req["stop"] = stop + if frequency_penalty != 0.0: + req["frequency_penalty"] = frequency_penalty + if presence_penalty != 0.0: + req["presence_penalty"] = presence_penalty + if seed is not None: + req["seed"] = seed + if reasoning_effort is not None: + req["reasoning_effort"] = reasoning_effort + if response_format is not None: + req["response_format"] = response_format + if tools is not None: + req["tools"] = tools + if tool_choice is not None: + req["tool_choice"] = tool_choice + if user is not None: + req["user"] = user + + t0 = _time.time() + response = client.chat.completions.create(**req) + elapsed = _time.time() - t0 + + choice = response.choices[0] + text = choice.message.content or "" + usage = response.usage or {} + + usage_dict: dict = {} + for attr in ("prompt_tokens", "completion_tokens", "total_tokens"): + if hasattr(usage, attr): + usage_dict[attr] = getattr(usage, attr) + if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details: + ctd = usage.completion_tokens_details + usage_dict["reasoning_tokens"] = getattr(ctd, "reasoning_tokens", None) + usage_dict["accepted_prediction_tokens"] = getattr(ctd, "accepted_prediction_tokens", None) + usage_dict["rejected_prediction_tokens"] = getattr(ctd, "rejected_prediction_tokens", None) + if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details: + ptd = usage.prompt_tokens_details + usage_dict["cached_tokens"] = getattr(ptd, "cached_tokens", None) + + tool_calls = None + if choice.message.tool_calls: + tool_calls = [ + { + "id": tc.id, + "type": tc.type, + "function": {"name": tc.function.name, "arguments": tc.function.arguments}, + } + for tc in choice.message.tool_calls + ] + + record_llm_call( + { + "type": "chat_completion", + "model": model, + "model_used": response.model, + "messages": messages, + "user_prompt": prompt, + "system_prompt": system, + "temperature": temperature, + "temperature_suppressed": temperature is None, + "top_p": top_p, + "max_tokens": max_tokens, + "n": n, + "stop": stop, + "frequency_penalty": frequency_penalty, + "presence_penalty": presence_penalty, + "seed": seed, + "reasoning_effort": reasoning_effort, + "response_format": response_format, + "tools_provided": [t.get("function", {}).get("name") for t in (tools or [])], + "tool_choice": tool_choice, + "text": text, + "finish_reason": choice.finish_reason, + "tool_calls": tool_calls, + "system_fingerprint": getattr(response, "system_fingerprint", None), + "response_id": response.id, + "created": response.created, + "usage": usage_dict, + "elapsed_s": elapsed, + "context": context or {}, + } + ) + return text + + +def anthropic_chat( + prompt: str, + model: str = "claude-haiku-4-5-20251001", + system: str = "You are a helpful assistant.", + max_tokens: int = 1024, + temperature: float | None = 1.0, + top_p: float | None = None, + top_k: int | None = None, + stop_sequences: list[str] | None = None, + tools: list | None = None, + tool_choice: dict | None = None, + thinking: dict | None = None, + metadata: dict | None = None, + context: dict | None = None, +) -> str: + """ + Make an Anthropic (Claude) chat completion call and record it for FlowCept provenance. + + Captures all request/response fields — including thinking blocks, tool use, + and cache token counts — as a child TaskObject. + """ + import anthropic as _anthropic + import os as _os + import time as _time + + client = _anthropic.Anthropic(api_key=_os.environ.get("ANTHROPIC_API_KEY")) + req: dict = { + "model": model, + "max_tokens": max_tokens, + "system": system, + "messages": [{"role": "user", "content": prompt}], + } + if temperature is not None: + req["temperature"] = temperature + if top_p is not None: + req["top_p"] = top_p + if top_k is not None: + req["top_k"] = top_k + if stop_sequences: + req["stop_sequences"] = stop_sequences + if tools: + req["tools"] = tools + if tool_choice: + req["tool_choice"] = tool_choice + if thinking: + req["thinking"] = thinking + if metadata: + req["metadata"] = metadata + + t0 = _time.time() + response = client.messages.create(**req) + elapsed = _time.time() - t0 + + text = "" + thinking_text = "" + tool_uses = [] + for block in response.content: + if block.type == "text": + text += block.text + elif block.type == "thinking": + thinking_text += getattr(block, "thinking", "") + elif block.type == "tool_use": + tool_uses.append({"id": block.id, "name": block.name, "input": block.input}) + + usage = response.usage + usage_dict = { + "input_tokens": getattr(usage, "input_tokens", None), + "output_tokens": getattr(usage, "output_tokens", None), + "cache_creation_input_tokens": getattr(usage, "cache_creation_input_tokens", None), + "cache_read_input_tokens": getattr(usage, "cache_read_input_tokens", None), + } + + record_llm_call( + { + "type": "chat_completion", + "model": model, + "model_used": response.model, + "messages": req["messages"], + "user_prompt": prompt, + "system_prompt": system, + "max_tokens": max_tokens, + "temperature": temperature, + "temperature_suppressed": temperature is None, + "top_p": top_p, + "top_k": top_k, + "stop_sequences": stop_sequences, + "thinking_budget_tokens": (thinking or {}).get("budget_tokens"), + "tools_provided": [t.get("name") for t in (tools or [])], + "tool_choice": tool_choice, + "text": text, + "thinking_text": thinking_text if thinking_text else None, + "finish_reason": response.stop_reason, + "stop_sequence": response.stop_sequence, + "tool_uses": tool_uses if tool_uses else None, + "response_id": response.id, + "usage": usage_dict, + "elapsed_s": elapsed, + "context": context or {}, + } + ) + return text + + +# --------------------------------------------------------------------------- +# FlowceptAnthropicClient — wraps anthropic.Anthropic / AsyncAnthropic to +# capture full provenance for every messages.create / messages.stream call. +# --------------------------------------------------------------------------- + + +class FlowceptAnthropicClient: + """ + Wrap an ``anthropic.Anthropic`` (or ``AsyncAnthropic``) client for provenance capture. + + Records every ``messages.create`` / ``messages.stream`` call as a FlowCept + provenance record (subtype=llm_call) via ``record_llm_call()``. + + Usage:: + + import anthropic + from flowcept.agents.crewai.crewai_plugin import FlowceptAnthropicClient + + client = FlowceptAnthropicClient(anthropic.Anthropic(), agent_name="my-agent") + response = client.messages.create( + model="claude-haiku-4-5-20251001", + max_tokens=1024, + messages=[{"role": "user", "content": "Hello"}], + ) + """ + + def __init__(self, client, agent_name=None, context=None): + self._inner = client + self._agent_name = agent_name + self._context: dict = context or {} + self.messages = _FlowceptAnthropicMessages(client.messages, agent_name, self._context) + + def __getattr__(self, name): + """Delegate attribute access to the wrapped client.""" + return getattr(self._inner, name) + + +class _FlowceptAnthropicMessages: + def __init__(self, messages_resource, agent_name, context): + self._inner = messages_resource + self._agent_name = agent_name + self._context = context + + def __getattr__(self, name): + return getattr(self._inner, name) + + def create(self, **kwargs): + import time as _time + + t0 = _time.time() + result = self._inner.create(**kwargs) + self._record(kwargs, result, _time.time() - t0) + return result + + async def async_create(self, **kwargs): + import time as _time + + t0 = _time.time() + result = await self._inner.create(**kwargs) + self._record(kwargs, result, _time.time() - t0) + return result + + def stream(self, **kwargs): + import time as _time + + return _FlowceptAnthropicStream(self._inner.stream(**kwargs), kwargs, self._record, _time.time()) + + def _record(self, kwargs, result, elapsed): + try: + model = kwargs.get("model", "unknown") + text = "" + thinking_text = "" + tool_uses = [] + for block in getattr(result, "content", []): + btype = getattr(block, "type", None) + if btype == "text": + text += getattr(block, "text", "") + elif btype == "thinking": + thinking_text += getattr(block, "thinking", "") + elif btype == "tool_use": + tool_uses.append( + { + "id": getattr(block, "id", None), + "name": getattr(block, "name", None), + "input": getattr(block, "input", None), + } + ) + usage = getattr(result, "usage", None) + usage_dict = ( + { + attr: getattr(usage, attr, None) + for attr in ( + "input_tokens", + "output_tokens", + "cache_creation_input_tokens", + "cache_read_input_tokens", + ) + } + if usage + else {} + ) + ctx = dict(self._context) + if self._agent_name: + ctx["agent_name"] = self._agent_name + payload = { + "type": "chat_completion", + "model": model, + "model_used": getattr(result, "model", model), + "messages": kwargs.get("messages"), + "system_prompt": kwargs.get("system"), + "max_tokens": kwargs.get("max_tokens"), + "temperature": kwargs.get("temperature"), + "top_p": kwargs.get("top_p"), + "top_k": kwargs.get("top_k"), + "stop_sequences": kwargs.get("stop_sequences"), + "thinking_budget_tokens": (kwargs.get("thinking") or {}).get("budget_tokens"), + "tools_provided": [t.get("name") for t in (kwargs.get("tools") or [])], + "tool_choice": kwargs.get("tool_choice"), + "text": text, + "thinking_text": thinking_text or None, + "finish_reason": getattr(result, "stop_reason", None), + "stop_sequence": getattr(result, "stop_sequence", None), + "tool_uses": tool_uses or None, + "response_id": getattr(result, "id", None), + "usage": usage_dict, + "elapsed_s": elapsed, + "context": ctx, + } + record_llm_call({k: v for k, v in payload.items() if v is not None}) + except Exception: + pass + + +class _FlowceptAnthropicStream: + def __init__(self, ctx_mgr, kwargs, record_fn, t0): + self._ctx_mgr = ctx_mgr + self._kwargs = kwargs + self._record_fn = record_fn + self._t0 = t0 + self._stream = None + + def __enter__(self): + self._stream = self._ctx_mgr.__enter__() + return self._stream + + def __exit__(self, *args): + import time as _time + + result = None + try: + result = self._stream.get_final_message() + except Exception: + pass + if result is not None: + self._record_fn(self._kwargs, result, _time.time() - self._t0) + return self._ctx_mgr.__exit__(*args) + + +# --------------------------------------------------------------------------- +# Public plugin class +# --------------------------------------------------------------------------- + + +class FlowceptCrewAIPlugin: + """ + FlowCept provenance plugin for CrewAI workflows. + + Registers a listener on CrewAI's native event bus so every crew kickoff, + task, agent execution, LLM call, and tool usage is captured automatically. + + Parameters + ---------- + config : dict, optional + Plugin configuration keys: + enabled (bool, default True) + workflow_name (str, default "crewai-workflow") + performance_tracking (bool, default True) + perf_csv (str, optional) — explicit path for timing CSV. + + Usage + ----- + Standalone:: + + plugin = FlowceptCrewAIPlugin(config={"workflow_name": "my-crew"}) + plugin.start() + crew.kickoff() + plugin.stop() + + Shared with Academy plugin:: + + academy_plugin = FlowceptAcademyPlugin(config={...}).start() + crewai_plugin = FlowceptCrewAIPlugin.from_academy_plugin(academy_plugin) + crew.kickoff() + academy_plugin.stop() # flushes the shared buffer + """ + + def __init__(self, config: dict | None = None, _shared_interceptor=None) -> None: + cfg = config or {} + self._enabled: bool = cfg.get("enabled", True) + self._workflow_name: str = cfg.get("workflow_name", "crewai-workflow") + self._campaign_id: str | None = cfg.get("campaign_id", None) + self._perf_tracking: bool = cfg.get("performance_tracking", True) + self._perf_csv: str | None = cfg.get("perf_csv", None) + self._shared_interceptor = _shared_interceptor + self._interceptor = _shared_interceptor or _CrewAIInterceptor() + self._owns_interceptor: bool = _shared_interceptor is None + self._stats: _ProvenanceStats | None = None + self._listener_obj: _FlowceptCrewAIListener | None = None + self._hooks_obj: _FlowceptCrewAIHooks | None = None + self._started = False + + def _register_hooks(self) -> None: + from crewai.hooks.llm_hooks import ( + register_before_llm_call_hook, + register_after_llm_call_hook, + ) + from crewai.hooks.tool_hooks import ( + register_before_tool_call_hook, + register_after_tool_call_hook, + ) + + register_before_llm_call_hook(self._hooks_obj.before_llm_call) + register_after_llm_call_hook(self._hooks_obj.after_llm_call) + register_before_tool_call_hook(self._hooks_obj.before_tool_call) + register_after_tool_call_hook(self._hooks_obj.after_tool_call) + + def _unregister_hooks(self) -> None: + if self._hooks_obj is None: + return + try: + from crewai.hooks.llm_hooks import ( + unregister_before_llm_call_hook, + unregister_after_llm_call_hook, + ) + from crewai.hooks.tool_hooks import ( + unregister_before_tool_call_hook, + unregister_after_tool_call_hook, + ) + + unregister_before_llm_call_hook(self._hooks_obj.before_llm_call) + unregister_after_llm_call_hook(self._hooks_obj.after_llm_call) + unregister_before_tool_call_hook(self._hooks_obj.before_tool_call) + unregister_after_tool_call_hook(self._hooks_obj.after_tool_call) + except Exception as e: + _log.debug("Hook unregister warning: %r", e) + + @classmethod + def from_academy_plugin( + cls, + academy_plugin: Any, + config: dict | None = None, + ) -> "FlowceptCrewAIPlugin": + """ + Create a CrewAI plugin that shares the buffer of a running FlowceptAcademyPlugin. + + All provenance records land in the same JSONL file when the Academy + plugin stops. + """ + interceptor = academy_plugin._interceptor + inst = cls(config=config, _shared_interceptor=interceptor) + inst._started = True + global _ACTIVE_INTERCEPTOR + _ACTIVE_INTERCEPTOR = interceptor + + inst._stats = _ProvenanceStats() if (config or {}).get("performance_tracking", True) else None + listener_obj, _ = _build_listener_class(interceptor, inst._stats) + inst._listener_obj = listener_obj + inst._hooks_obj = _FlowceptCrewAIHooks(interceptor, inst._stats, listener_obj) + inst._register_hooks() + global _PROV_STATS + _PROV_STATS = inst._stats + return inst + + def start(self) -> "FlowceptCrewAIPlugin": + """Start provenance capture: register listeners, hooks, and the interceptor.""" + if not self._enabled or self._started: + return self + if not self._owns_interceptor: + return self + try: + self._stats = _ProvenanceStats() if self._perf_tracking else None + global _PROV_STATS + _PROV_STATS = self._stats + self._interceptor.start(self._workflow_name, campaign_id=self._campaign_id) + self._campaign_id = self._interceptor._campaign_id + listener_obj, _ = _build_listener_class(self._interceptor, self._stats) + self._listener_obj = listener_obj + self._hooks_obj = _FlowceptCrewAIHooks(self._interceptor, self._stats, listener_obj) + self._register_hooks() + self._started = True + global _ACTIVE_INTERCEPTOR + _ACTIVE_INTERCEPTOR = self._interceptor + wf_id = self._interceptor._workflow_id + print( + f"[FlowceptCrewAIPlugin] Started\n" + f" workflow_id : {wf_id}\n" + f" campaign_id : {self._interceptor._campaign_id}\n" + f" Capturing : crew/task/agent (event bus) + " + f"LLM/tool (hooks — full messages, response, agent context).", + flush=True, + ) + except Exception as e: + print( + f"[FlowceptCrewAIPlugin] WARNING: failed to start — {e!r}. Continuing without provenance capture.", + flush=True, + ) + _log.exception("FlowceptCrewAIPlugin start failed") + self._enabled = False + return self + + def stop(self) -> None: + """Stop provenance capture: unregister hooks and flush records.""" + if not self._started: + return + self._unregister_hooks() + if not self._owns_interceptor: + self._started = False + print( + "[FlowceptCrewAIPlugin] Detached from shared buffer (flushed by the owning plugin).", + flush=True, + ) + self._maybe_write_perf_csv() + global _PROV_STATS + _PROV_STATS = None + return + try: + _t0 = time.perf_counter() + self._interceptor.stop() + if self._stats is not None: + self._stats.record("flush", time.perf_counter() - _t0) + except Exception as e: + print(f"[FlowceptCrewAIPlugin] Warning during stop: {e!r}", flush=True) + self._started = False + print("[FlowceptCrewAIPlugin] Stopped.", flush=True) + self._maybe_write_perf_csv() + global _ACTIVE_INTERCEPTOR + _ACTIVE_INTERCEPTOR = None + _PROV_STATS = None + + def _maybe_write_perf_csv(self) -> None: + if self._stats is None: + return + print( + "\n[FlowceptCrewAIPlugin] Provenance overhead report:\n" + + self._stats.summary() + + "\n (N = event count; Total/Mean/Min/Max in ms/µs respectively)\n", + flush=True, + ) + wf_id = self._interceptor._workflow_id + csv_path = self._perf_csv or f"crewai_provenance_perf_{wf_id}.csv" + try: + self._stats.to_csv(csv_path, workflow_id=wf_id) + print(f"[FlowceptCrewAIPlugin] Performance stats → {csv_path}", flush=True) + except Exception as e: + print(f"[FlowceptCrewAIPlugin] Warning: could not write perf CSV — {e!r}", flush=True) + + def __enter__(self) -> "FlowceptCrewAIPlugin": + """Start the plugin when entering a context manager block.""" + return self.start() + + def __exit__(self, *_: Any) -> None: + """Stop the plugin when exiting a context manager block.""" + self.stop() diff --git a/src/flowcept/agents/harness/README.md b/src/flowcept/agents/harness/README.md new file mode 100644 index 00000000..60134950 --- /dev/null +++ b/src/flowcept/agents/harness/README.md @@ -0,0 +1,318 @@ +# AI coding harness provenance (`flowcept.agents.harness`) + +Capture what an AI coding harness actually did — prompts, turns, tool calls, +subagents — as [Flowcept](https://github.com/ORNL/flowcept) provenance. + +This package is the shared capture core; the per-harness plugin modules live +beside it under `flowcept/agents/` (`claude_code/`, `cli_harness/`, `otel/`, +`claude_agent_sdk/`, `openai_agents/`, `langchain/`), following the same layout +as the agent-framework plugins (`academy/`, `langgraph/`, `crewai/`, +`autogen/`). + +Agentic coding sessions are workflows: a prompt causes a turn, a turn causes +tool calls, a tool call edits a file. That is exactly the structure Flowcept +already stores and queries for scientific workflows, and +[PROV-AGENT](https://arxiv.org/abs/2508.02866) is the W3C PROV extension that +names the pieces. This package writes that structure out of the harnesses you +already use, so "which prompt produced this bad edit?" is a query rather than a +scroll back through a transcript. + +Supported sources: + +| Source | How | +| --- | --- | +| Claude Code | plugin (hooks), or hooks in `settings.json` | +| Codex CLI, Gemini CLI, Cursor, OpenCode | generic hook adapter + a JSON profile | +| Anything emitting OpenTelemetry GenAI spans | span exporter, or ingest exported spans | +| Claude Agent SDK | `trace_query`, or feed the message stream to a tracer | +| OpenAI Agents SDK | a tracing processor | +| LangChain / LangGraph | a callback handler | +| Your own agent | `SessionTracer`, or the MCP server's `record_event` tool | + +## Install + +Ships with flowcept itself: + +```bash +pip install flowcept # capture, query, and report +pip install "flowcept[harness_otel]" # + the OTel span exporter +pip install "flowcept[harness_claude_sdk]" # + the Claude Agent SDK wrapper +``` + +The capture path is deliberately stdlib-only — it imports nothing outside +`flowcept.agents.harness` and the standard library. A Claude Code hook is a +fresh process on the interactive critical path, and importing flowcept's heavy +dependencies costs far more than the capture itself; those are only loaded to +*read* what was captured. + +## Quick start: Claude Code + +``` +/plugin marketplace add +/plugin install flowcept +``` + +Then work normally. When you want to see what was recorded: + +```bash +flowcept-harness sessions # every captured session, newest first +flowcept-harness show # the most recent one, turn by turn +flowcept-harness report # a Flowcept workflow card +``` + +`flowcept-harness install` prints the equivalent `settings.json` if you would +rather wire the hooks yourself than use the plugin. + +## Quick start: another CLI harness + +Point the harness's hook at the generic adapter with the matching profile: + +```bash +flowcept-harness hook --harness codex --profile codex +``` + +Profiles live in [`src/flowcept/agents/cli_harness/profiles/`](src/flowcept/agents/cli_harness/profiles/) +and are plain JSON: a map from the harness's event names to normalized ones, +and a map from its payload fields to ours. Adding a harness means adding a file, +not writing code. `--profile` also accepts a path, so a profile can live outside +the package while you iterate on it. + +## Quick start: OpenTelemetry + +In-process, for anything already instrumented: + +```python +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from flowcept.agents.otel.otel_plugin import FlowceptSpanExporter + +provider = TracerProvider() +provider.add_span_processor(SimpleSpanProcessor(FlowceptSpanExporter())) +``` + +Or after the fact, from spans a collector already wrote: + +```bash +python -c "from flowcept.agents.otel.otel_plugin import ingest_file; ingest_file('spans.jsonl')" +``` + +Spans are read through the OTel GenAI semantic conventions: +`gen_ai.operation.name` separates a model call from a tool call, +`gen_ai.conversation.id` groups spans into a session. Spans that are not GenAI +spans are ignored — an HTTP client span is not provenance. + +## Quick start: the SDKs + +```python +# Claude Agent SDK — a drop-in for claude_agent_sdk.query +from flowcept.agents.claude_agent_sdk.claude_agent_sdk_plugin import trace_query + +async for message in trace_query(prompt="fix the failing test"): + ... + +# OpenAI Agents SDK — register once, nothing else changes +from flowcept.agents.openai_agents.openai_agents_plugin import install + +install() + +# LangChain / LangGraph — pass the handler as a callback +from flowcept.agents.langchain.langchain_plugin import FlowceptCallbackHandler + +graph.invoke(state, config={"callbacks": [FlowceptCallbackHandler(session_id="thread-42")]}) +``` + +Each wrapper is duck-typed against its SDK — nothing imports the SDK it wraps, +so installing one does not drag in the others. + +For an agent that is none of the above, drive the session yourself: + +```python +from flowcept.agents.harness import SessionTracer + +with SessionTracer("my_agent", model="claude-opus-5") as tracer: + tracer.prompt("summarize the repo") + with tracer.tool("read_file", {"path": "README.md"}) as call: + call.result(read("README.md")) + tracer.turn_end("done", usage={"input_tokens": 900}) +``` + +## What gets recorded + +``` +session workflow (subtype: agent_session) + turn task (subtype: ai_model_invocation, granularity=turn) + model call task (subtype: ai_model_invocation, granularity=call) + tool call task (subtype: agent_tool) + subagent workflow (subtype: subagent_session) + task + compaction, notice task (subtype: harness_event) +``` + +The edges are the point. A tool call's `parent_task_id` is its turn; a +subagent's tools live in the subagent's own workflow rather than interleaved +with the parent's; a turn lists the tools it caused. That is the +`wasInformedBy` chain PROV-AGENT is built around, and it is what lets you walk +from a bad edit back to the prompt that caused it. + +Model invocations are recorded at whatever granularity the source can see. A +hook cannot observe individual API calls, so hook-based capture records one +invocation per turn; SDK and OTel capture record both. `granularity` in +`custom_metadata` says which you are looking at. + +### Linking to other capture systems + +Harness sessions can point back at provenance emitted by the agent-framework +plugins (LangGraph, Academy, AutoGen, ...). Give the harness a +framework-emitted task or agent id and every turn, tool, and LLM-call task it +records carries it as `source_agent_id`: + +- **hook payload** — a `flowcept_source_agent_id` key in the payload (set by a + wrapper or whatever launches the harness) wins for that event; +- **environment** — `FLOWCEPT_HARNESS_SOURCE_AGENT_ID` applies to every event + of the process it is set for; +- **SessionTracer** — pass `source_agent_id=...` when constructing the tracer. + +The reverse direction (a framework record pointing at a harness task) is the +framework plugins' `_source_agent_id` / `source_agent_id` input, and +`flowcept-harness analyze --links` walks both kinds of edge. + +## Where it goes + +Records are appended as JSONL, one file per session, under +`~/.flowcept/harness/buffers/`. The format is Flowcept's own, so the buffer is +directly consumable: + +```bash +flowcept --generate-report --input-path ~/.flowcept/harness/buffers/.jsonl +``` + +To push into a live Flowcept backend instead of (or as well as) the file: + +```bash +export FLOWCEPT_HARNESS_ONLINE=1 # publish as you go +flowcept-harness flush --all # or publish buffers after the fact +``` + +Offline is the default because a hook must never block on a message queue that +may not be running. + +## Configuration + +Every knob is an environment variable, so it can be set from a harness settings +file, a plugin's `userConfig`, or a shell profile. + +| Variable | Default | Meaning | +| --- | --- | --- | +| `FLOWCEPT_HARNESS_ENABLED` | `1` | Master switch. | +| `FLOWCEPT_HARNESS_HOME` | `~/.flowcept/harness` | State and buffers. | +| `FLOWCEPT_HARNESS_BUFFER_DIR` | *(under home)* | Buffers elsewhere. | +| `FLOWCEPT_HARNESS_CAMPAIGN_SCOPE` | `project` | `project`, `global`, or `none`. | +| `FLOWCEPT_HARNESS_CAMPAIGN_ID` | *(derived)* | Pin sessions to one campaign. | +| `FLOWCEPT_HARNESS_CONTENT` | `summary` | File bodies: `full`, `summary`, `none`. | +| `FLOWCEPT_HARNESS_MAX_STR` | `4000` | Max characters per captured string. | +| `FLOWCEPT_HARNESS_REDACT` | `1` | Redact credential-shaped keys and literals. | +| `FLOWCEPT_HARNESS_CAPTURE_PROMPTS` | `1` | Off stores prompt digests only. | +| `FLOWCEPT_HARNESS_CAPTURE_TOOL_RESULTS` | `1` | Off stores inputs but not outputs. | +| `FLOWCEPT_HARNESS_SOURCE_AGENT_ID` | *(unset)* | Link tasks to another system's task/agent id. | +| `FLOWCEPT_HARNESS_ONLINE` | `0` | Publish to the Flowcept MQ as you go. | +| `FLOWCEPT_HARNESS_TIMEOUT_MS` | `2000` | Hard ceiling on hook wall time. | +| `FLOWCEPT_HARNESS_DEBUG` | `0` | Log capture failures instead of staying silent. | + +### Privacy + +Prompts and tool inputs are captured by default, because provenance without +them answers very little. What is *not* captured: values under +credential-shaped keys and literals matching known key formats are replaced +with `«redacted»` at capture time, not at query time, since provenance is +long-lived and often shared. File bodies in `Write`/`Edit` inputs are reduced to +a size, a line count, a hash, and a 200-character preview. + +For a stricter posture, `FLOWCEPT_HARNESS_CONTENT=none` drops file bodies +entirely and `FLOWCEPT_HARNESS_CAPTURE_PROMPTS=0` keeps only a digest of each +prompt — enough to tell two prompts apart, not enough to read them. + +## The MCP server + +Exposes captured provenance to an agent as tools, so a session can ask about +its own history: + +```bash +flowcept-harness-mcp +``` + +Tools: `list_sessions`, `get_session`, `search_tool_calls`, `session_stats`, +`record_event`, `generate_report`. `record_event` also makes the server a +capture path in its own right, for a harness that speaks MCP but has no hooks. + +## Never harming the harness + +Capture code runs inside an interactive tool, which constrains it more than +correctness alone would: + +- **Nothing on stdout.** Claude Code injects hook stdout into the model's + context on some events; a provenance record must not become a prompt. +- **Always exit 0.** A capture failure is logged, never reported to the user as + a broken hook. +- **A watchdog.** Past `TIMEOUT_MS` the process hard-exits. Losing one record + beats stalling the UI. +- **Crash-safe records.** A session's workflow record is written when it opens, + so an interrupted run is still readable, and rewritten when it closes. +- **Concurrency-safe.** Session state is a locked read-modify-write and + appends are locked, so parallel subagents cannot interleave a record. + +`flowcept-harness repair` closes sessions left open by a harness that died. + +## CLI + +``` +flowcept-harness sessions list captured sessions +flowcept-harness show show one session's activity +flowcept-harness analyze analyze a session (--errors, --slowest, --links, --compare A B) +flowcept-harness status configuration and capture health +flowcept-harness report generate a Flowcept report +flowcept-harness flush publish buffers to a Flowcept backend +flowcept-harness repair close sessions a crashed harness left open +flowcept-harness install print the settings that enable capture +flowcept-harness hook record a payload from stdin (what hooks call) +``` + +## Development + +From the repository root: + +```bash +python -m venv .venv +.venv/bin/pip install -e ".[dev]" "mcp>=1.0.0" "opentelemetry-sdk>=1.20.0" +.venv/bin/python -m pytest tests/harness +``` + +Tests drive the real thing wherever there is one: the Flowcept interop tests +run against installed flowcept, the MCP tests run a real stdio client +handshake, and the OTel exporter test runs through a real tracer provider. The +SDK wrappers are tested against payload-shaped objects, since duck-typing is +the contract they are written to. + +## Layout + +``` +src/flowcept/agents/ + harness/ the shared capture core (this package) + events.py the harness-independent event every adapter produces + recorder.py the state machine that turns events into PROV-AGENT records + prov.py record constructors + ids.py deterministic UUIDv5 ids, so separate processes agree + emit.py JSONL buffers, locking, and online publishing + sanitize.py redaction, truncation, JSON-safety + state.py locked per-session state + tracer.py SessionTracer, for agents you write yourself + cli.py the command line + mcp_server.py provenance as MCP tools + claude_code/ Claude Code hook adapter + cli_harness/ profile-driven adapter (+ profiles/: codex, gemini, cursor, opencode) + otel/ OTel GenAI span exporter and ingest + claude_agent_sdk/ trace_query wrapper + openai_agents/ tracing processor + langchain/ callback handler +plugins/flowcept/ the Claude Code plugin +tests/harness/ the test suite +examples/agents/harness/ a SessionTracer example +``` diff --git a/src/flowcept/agents/harness/__init__.py b/src/flowcept/agents/harness/__init__.py new file mode 100644 index 00000000..7a75bc38 --- /dev/null +++ b/src/flowcept/agents/harness/__init__.py @@ -0,0 +1,36 @@ +"""flowcept-harness: PROV-AGENT provenance capture for AI coding harnesses. + +Captures what an agentic harness actually did — prompts, turns, tool calls, +subagents — as `Flowcept `_ provenance, +modelled with PROV-AGENT (arXiv:2508.02866). + +The capture path is stdlib-only and writes JSONL that Flowcept reads natively; +flowcept itself is only needed to ingest or query what was captured. + + from flowcept.agents.harness import HarnessEvent, Recorder, EventKind + + Recorder().record(HarnessEvent( + kind=EventKind.TOOL_POST, + harness="my_harness", + session_id="abc123", + tool_name="run_tests", + tool_input={"suite": "unit"}, + tool_response={"passed": 42}, + )) +""" + +from .config import Config, load_config +from .events import HarnessEvent +from .recorder import Recorder, repair_session +from .tracer import SessionTracer +from .vocab import EventKind + +__all__ = [ + "Config", + "EventKind", + "HarnessEvent", + "Recorder", + "SessionTracer", + "load_config", + "repair_session", +] diff --git a/src/flowcept/agents/harness/cli.py b/src/flowcept/agents/harness/cli.py new file mode 100644 index 00000000..cf1e816d --- /dev/null +++ b/src/flowcept/agents/harness/cli.py @@ -0,0 +1,527 @@ +"""The ``flowcept-harness`` command. + +Everything outside the capture path lives here: inspecting what was captured, +pushing it into Flowcept, and wiring harnesses up in the first place. Unlike +the hook path this is not latency-sensitive, so it imports freely. +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +from collections.abc import Iterator +from pathlib import Path +from typing import Any + +from flowcept.version import __version__ + +from .config import Config, load_config + +# Exit codes. Anything non-zero means the *command* failed; capture failures +# are always silent by design. +OK = 0 +FAILED = 1 + + +# -- helpers ----------------------------------------------------------------- + + +def _buffers(config: Config) -> list[Path]: + if not config.buffers_dir.is_dir(): + return [] + return sorted(config.buffers_dir.glob("*.jsonl"), key=lambda p: p.stat().st_mtime, reverse=True) + + +def _read_records(path: Path) -> Iterator[dict[str, Any]]: + with path.open("r", encoding="utf-8") as fh: + for line in fh: + line = line.strip() + if not line: + continue + try: + record = json.loads(line) + except ValueError: + continue + if isinstance(record, dict): + yield record + + +def _resolve_inputs(config: Config, given: list[str] | None, want_all: bool) -> list[Path]: + if given: + return [Path(p).expanduser() for p in given] + if want_all: + return _buffers(config) + latest = _buffers(config) + return latest[:1] + + +def _summarize(path: Path) -> dict[str, Any]: + counts = {"workflow": 0, "task": 0, "agent": 0, "other": 0} + session: dict[str, Any] = {} + for record in _read_records(path): + kind = record.get("type") + counts[kind if kind in counts else "other"] += 1 + if kind == "workflow" and record.get("parent_workflow_id") is None: + session = record + return { + "path": path, + "counts": counts, + "workflow_id": session.get("workflow_id") or path.stem, + "status": session.get("status") or "UNKNOWN", + "name": session.get("name") or "?", + "started_at": session.get("started_at"), + "ended_at": session.get("ended_at"), + "generated": session.get("generated") or {}, + } + + +def _fmt_time(value: Any) -> str: + if not isinstance(value, (int, float)): + return "-" + import datetime + + return datetime.datetime.fromtimestamp(value).strftime("%Y-%m-%d %H:%M") + + +# -- commands ---------------------------------------------------------------- + + +def cmd_hook(args, config: Config) -> int: + """Route a hook payload from stdin to the right adapter. + + Lets one binary serve every harness: ``flowcept-harness hook --harness + codex`` is a valid hook command anywhere a program can be named. + """ + from .runtime import read_stdin_json, run_capture + + def _run(cfg: Config) -> None: + payload = read_stdin_json() + if args.harness == "claude_code": + from flowcept.agents.claude_code import claude_code_plugin as claude_code + + if args.event: + payload["hook_event_name"] = payload.get("hook_event_name") or args.event + claude_code.handle(payload, cfg) + else: + from flowcept.agents.cli_harness import cli_harness_plugin as generic + + generic.handle(payload, cfg, harness=args.harness, profile=args.profile, event=args.event) + + return run_capture(_run, config) + + +def cmd_status(args, config: Config) -> int: + """Print configuration, capture health, and optionally probe the backend.""" + buffers = _buffers(config) + print(f"flowcept-harness {__version__}") + print(f" enabled: {config.enabled}") + print(f" home: {config.home}") + print(f" buffers: {config.buffers_dir} ({len(buffers)} session(s))") + print(f" content mode: {config.content_mode} redact: {config.redact}") + print(f" online: {config.online}") + + try: + import flowcept + + print(f" flowcept: {getattr(flowcept, '__version__', 'installed')}") + except ImportError: + print(" flowcept: not installed (capture works; ingest and reports do not)") + return OK + + if args.check_backend: + try: + from flowcept.commons.daos.mq_dao.mq_dao_base import MQDao + + MQDao.build() + print(" backend: reachable") + except Exception as exc: + print(f" backend: unreachable ({type(exc).__name__}: {exc})") + return OK + + +def cmd_sessions(args, config: Config) -> int: + """List captured sessions, newest first.""" + buffers = _buffers(config) + if not buffers: + print(f"No sessions captured yet under {config.buffers_dir}") + return OK + + for path in buffers[: args.limit]: + info = _summarize(path) + gen = info["generated"] + detail = " ".join(f"{k}={v}" for k, v in gen.items()) or "-" + print(f"{info['workflow_id'][:8]} {info['status']:8} {_fmt_time(info['started_at'])} {info['name']}") + print(f" {detail}") + if args.verbose: + counts = info["counts"] + print(f" records: {counts} file: {path}") + return OK + + +def cmd_show(args, config: Config) -> int: + """Show the recorded activity of one session, with subagent work indented.""" + paths = _resolve_inputs(config, args.input, want_all=False) + if not paths: + print("No session to show.", file=sys.stderr) + return FAILED + + path = paths[0] + records = list(_read_records(path)) + # Subagent work is indented under the session it belongs to. + nested = { + r["workflow_id"] for r in records if r.get("type") == "workflow" and r.get("parent_workflow_id") is not None + } + + tasks = [r for r in records if r.get("type") == "task"] + if not tasks: + print(f"{path.name}: no activity recorded yet.") + return OK + + for record in tasks: + indent = " " if record.get("workflow_id") in nested else " " + elapsed = "" + started, ended = record.get("started_at"), record.get("ended_at") + if isinstance(started, (int, float)) and isinstance(ended, (int, float)): + elapsed = f"{ended - started:.2f}s" + print( + f"{indent}{record.get('subtype', ''):20} {record.get('activity_id', '?'):24} " + f"{record.get('status', ''):9} {elapsed:>8}" + ) + if args.verbose and record.get("stderr"): + print(f"{indent} ! {record['stderr']}") + return OK + + +def cmd_flush(args, config: Config) -> int: + """Publish buffered records into a running Flowcept backend.""" + try: + from flowcept.commons.daos.mq_dao.mq_dao_base import MQDao + except ImportError: + print("flowcept is not installed; `pip install flowcept-harness[ingest]`", file=sys.stderr) + return FAILED + + paths = _resolve_inputs(config, args.input, args.all) + if not paths: + print("Nothing to flush.", file=sys.stderr) + return FAILED + + try: + mq = MQDao.build() + except Exception as exc: + print(f"Cannot reach the Flowcept backend: {exc}", file=sys.stderr) + return FAILED + + total = 0 + for path in paths: + records = list(_read_records(path)) + if not records: + continue + if args.dry_run: + print(f"would publish {len(records):5} records from {path.name}") + else: + mq.bulk_publish(records) + print(f"published {len(records):5} records from {path.name}") + total += len(records) + + if not args.dry_run: + try: + # check_safe_stops=False: a flush has no interceptor instance to + # coordinate, and the control messages it would send carry a None + # id that crashes the document inserter's bookkeeping. + mq.stop(check_safe_stops=False) + except Exception: + # Best effort: the records are already published. + pass + if args.remove: + for path in paths: + path.unlink(missing_ok=True) + path.with_suffix(path.suffix + ".lock").unlink(missing_ok=True) + + print(f"{'would publish' if args.dry_run else 'published'} {total} records from {len(paths)} file(s)") + return OK + + +def cmd_report(args, config: Config) -> int: + """Generate a Flowcept report from a captured buffer file.""" + try: + from flowcept import Flowcept + except ImportError: + print("flowcept is not installed; `pip install flowcept-harness[ingest]`", file=sys.stderr) + return FAILED + + paths = _resolve_inputs(config, args.input, want_all=False) + if not paths: + print("No session to report on.", file=sys.stderr) + return FAILED + + stats = Flowcept.generate_report( + report_type=args.type, + input_jsonl_path=str(paths[0]), + format=args.format, + output_path=args.output, + ) + if args.output: + print(f"Wrote {args.output} ({stats.get('n_workflows')} workflow(s), {stats.get('n_tasks')} task(s))") + else: + print(stats.get("markdown") or stats) + return OK + + +def _find_session_buffer(config: Config, session: str | None) -> Path | None: + buffers = _buffers(config) + if not buffers: + return None + if not session: + return buffers[0] + for path in buffers: + if path.stem == session or path.stem.startswith(session): + return path + return None + + +def _analyze_compare(args, config: Config, prov_core) -> int: + """Print a per-activity comparison of two captured sessions.""" + session_a, session_b = args.compare + paths = [] + for session in (session_a, session_b): + path = _find_session_buffer(config, session) + if path is None: + print(f"No session matching {session!r}.", file=sys.stderr) + return FAILED + paths.append(path) + path_a, path_b = paths + + comparison = prov_core.compare_executions(list(_read_records(path_a)), list(_read_records(path_b))) + totals = comparison["totals"] + + def _secs(value: Any) -> str: + return f"{value:.2f}s" if isinstance(value, (int, float)) else "?" + + def _rate(value: Any) -> str: + return f"{value:.0%}" if isinstance(value, (int, float)) else "?" + + print(f"comparing: A={path_a.stem} B={path_b.stem}") + tasks_delta = totals["n_tasks_b"] - totals["n_tasks_a"] + print(f"tasks: {totals['n_tasks_a']} -> {totals['n_tasks_b']} ({tasks_delta:+d})") + if totals["total_elapsed_delta"] is not None: + print( + f"elapsed: {_secs(totals['total_elapsed_a'])} -> {_secs(totals['total_elapsed_b'])} " + f"({totals['total_elapsed_delta']:+.2f}s)" + ) + for activity, row in comparison["activities"].items(): + line = f" {activity:24} count {row['count_a']} -> {row['count_b']} ({row['count_delta']:+d})" + if row["elapsed_avg_delta"] is not None: + avg_a, avg_b = _secs(row["elapsed_avg_a"]), _secs(row["elapsed_avg_b"]) + line += f" avg {avg_a} -> {avg_b} ({row['elapsed_avg_delta']:+.2f}s)" + if row["error_rate_a"] is not None or row["error_rate_b"] is not None: + line += f" errors {_rate(row['error_rate_a'])} -> {_rate(row['error_rate_b'])}" + print(line) + if comparison["only_in_a"]: + print(f"only in A: {', '.join(comparison['only_in_a'])}") + if comparison["only_in_b"]: + print(f"only in B: {', '.join(comparison['only_in_b'])}") + return OK + + +def cmd_analyze(args, config: Config) -> int: + """Analyze one captured session with the provenance analysis functions.""" + from flowcept.agents.prov_analysis import core as prov_core + + if args.compare: + # --compare drives its own two-session resolution; the single-session + # analyses make no sense alongside it. + if args.errors or args.links or args.slowest is not None: + print("--compare cannot be combined with --errors, --slowest, or --links.", file=sys.stderr) + return FAILED + return _analyze_compare(args, config, prov_core) + + path = _find_session_buffer(config, args.session) + if path is None: + print(f"No session matching {args.session!r}.", file=sys.stderr) + return FAILED + + records = list(_read_records(path)) + print(f"session: {path.stem}") + + if args.errors: + errors = prov_core.analyze_errors(records) + print(f"failed tasks: {errors['n_failed']} of {errors['n_tasks']}") + if errors["first_failure_at_utc"]: + print(f"first failure: {errors['first_failure_at_utc']} last: {errors['last_failure_at_utc']}") + for activity, entry in errors["by_activity"].items(): + rate = f"{entry['error_rate']:.0%}" if entry["error_rate"] is not None else "?" + print(f" {activity:24} {entry['n_failed']}/{entry['n_total']} failed ({rate})") + for excerpt in entry["excerpts"]: + print(f" ! {excerpt}") + return OK + + if args.slowest is not None: + for row in prov_core.find_slowest_tasks(records, limit=args.slowest): + print( + f" {row['elapsed_seconds']:>10.3f}s {str(row.get('activity_id') or '?'):24} " + f"{str(row.get('status') or ''):9} depth={row['parent_depth']}" + ) + return OK + + if args.links: + links = prov_core.cross_framework_links(records) + print(f"cross-framework links: {links['n_links']} unlinked tasks: {links['n_unlinked_tasks']}") + if links["frameworks_seen"]: + print(f"frameworks seen: {', '.join(links['frameworks_seen'])}") + for link in links["links"]: + frameworks = "<->".join(link["frameworks"]) or "?" + print(f" {link['source_task_id']} -> {link['target_task_id']} [{frameworks}]") + return OK + + summary = prov_core.summarize_execution(records) + behavior = prov_core.analyze_agent_behavior(records) + print(f"records: {summary['n_records']} workflows: {summary['n_workflows']} tasks: {summary['n_tasks']}") + if summary["total_elapsed_seconds"] is not None: + print(f"elapsed: {summary['total_elapsed_seconds']:.2f}s ({summary['started_at_utc']} UTC)") + print(f"statuses: {' '.join(f'{k}={v}' for k, v in summary['status_counts'].items()) or '-'}") + print(f"by subtype: {' '.join(f'{k}={v}' for k, v in summary['tasks_by_subtype'].items()) or '-'}") + usage = summary["token_usage"]["totals"] + if usage: + print(f"token usage: {' '.join(f'{k}={v}' for k, v in usage.items())}") + for session in behavior["sessions"]: + print( + f"session workflow: {session['workflow_id'][:8]} {session['status']} subagents={session['n_subagents']}" + ) + for agent, entry in behavior["agents"].items(): + tools = " ".join(f"{k}={v}" for k, v in entry["tool_calls_by_tool"].items()) or "-" + print(f" agent {agent[:8]}: turns={entry['turns']} llm_calls={entry['llm_calls']} tools: {tools}") + return OK + + +def cmd_repair(args, config: Config) -> int: + """Close sessions whose harness exited without a session-end event.""" + from .recorder import repair_session + + if not config.sessions_dir.is_dir(): + print("No sessions to repair.") + return OK + + repaired = 0 + for state_path in sorted(config.sessions_dir.glob("*.json")): + records = repair_session(config, state_path.stem) + if records: + repaired += 1 + print(f"repaired {state_path.stem[:8]} ({len(records)} record(s))") + print(f"{repaired} session(s) repaired.") + return OK + + +def cmd_install(args, config: Config) -> int: + """Print the settings needed to enable capture in a harness.""" + if args.harness == "claude_code": + print("Add the plugin marketplace, then enable the plugin:\n") + print(" /plugin marketplace add ") + print(" /plugin install flowcept\n") + print("Or wire the hooks directly in settings.json:\n") + events = [ + "SessionStart", + "SessionEnd", + "UserPromptSubmit", + "Stop", + "PreToolUse", + "PostToolUse", + "SubagentStart", + "SubagentStop", + ] + hooks = { + event: [{"hooks": [{"type": "command", "command": f"flowcept-harness hook --event {event}"}]}] + for event in events + } + print(json.dumps({"hooks": hooks}, indent=2)) + else: + print(f"Set the hook command for {args.harness} to:\n") + print(f" flowcept-harness hook --harness {args.harness}") + return OK + + +# -- parser ------------------------------------------------------------------ + + +def build_parser() -> argparse.ArgumentParser: + """Build the argument parser for the ``flowcept-harness`` command.""" + parser = argparse.ArgumentParser( + prog="flowcept-harness", + description="PROV-AGENT provenance capture for AI coding harnesses.", + ) + parser.add_argument("--version", action="version", version=f"flowcept-harness {__version__}") + parser.add_argument("--home", help="Override the capture home directory.") + sub = parser.add_subparsers(dest="command", required=True) + + p = sub.add_parser("hook", help="Record a hook payload read from stdin.") + p.add_argument("--harness", default="claude_code") + p.add_argument("--event", help="Event name, when the payload omits it.") + p.add_argument("--profile", help="Field-mapping profile for generic harnesses.") + p.set_defaults(func=cmd_hook) + + p = sub.add_parser("status", help="Show configuration and capture health.") + p.add_argument("--check-backend", action="store_true", help="Also probe the Flowcept backend.") + p.set_defaults(func=cmd_status) + + p = sub.add_parser("sessions", help="List captured sessions, newest first.") + p.add_argument("-n", "--limit", type=int, default=20) + p.add_argument("-v", "--verbose", action="store_true") + p.set_defaults(func=cmd_sessions) + + p = sub.add_parser("show", help="Show the activity of one session.") + p.add_argument("input", nargs="*", help="Buffer file (default: most recent).") + p.add_argument("-v", "--verbose", action="store_true") + p.set_defaults(func=cmd_show) + + p = sub.add_parser("flush", help="Publish buffered records to Flowcept.") + p.add_argument("--input", nargs="*", help="Buffer files (default: most recent).") + p.add_argument("--all", action="store_true", help="Flush every buffer.") + p.add_argument("--remove", action="store_true", help="Delete buffers after a successful flush.") + p.add_argument("--dry-run", action="store_true") + p.set_defaults(func=cmd_flush) + + p = sub.add_parser("report", help="Generate a Flowcept report from a buffer.") + p.add_argument("--input", nargs="*", help="Buffer file (default: most recent).") + p.add_argument("--type", default="workflow_card") + p.add_argument("--format", default="markdown") + p.add_argument("-o", "--output", help="Write to a file instead of stdout.") + p.set_defaults(func=cmd_report) + + p = sub.add_parser("analyze", help="Analyze one captured session's provenance.") + # A single session and a two-session comparison are different modes, so + # argparse rejects `analyze --compare A B` outright. + mode = p.add_mutually_exclusive_group() + mode.add_argument("session", nargs="?", help="Workflow id or prefix (default: most recent).") + mode.add_argument( + "--compare", + nargs=2, + metavar=("SESSION_A", "SESSION_B"), + help="Compare two sessions per activity (counts, durations, error rates).", + ) + p.add_argument("--errors", action="store_true", help="Analyze failures only.") + p.add_argument("--slowest", type=int, metavar="N", help="Show the N slowest tasks.") + p.add_argument("--links", action="store_true", help="Show cross-framework links.") + p.set_defaults(func=cmd_analyze) + + p = sub.add_parser("repair", help="Close sessions left open by a crashed harness.") + p.set_defaults(func=cmd_repair) + + p = sub.add_parser("install", help="Print the settings that enable capture.") + p.add_argument("--harness", default="claude_code") + p.set_defaults(func=cmd_install) + + return parser + + +def main(argv: list[str] | None = None) -> int: + """Parse arguments, load the configuration, and dispatch to a subcommand.""" + args = build_parser().parse_args(argv) + if args.home: + os.environ["FLOWCEPT_HARNESS_HOME"] = args.home + config = load_config() + return args.func(args, config) + + +if __name__ == "__main__": # pragma: no cover + raise SystemExit(main()) diff --git a/src/flowcept/agents/harness/config.py b/src/flowcept/agents/harness/config.py new file mode 100644 index 00000000..d9be539e --- /dev/null +++ b/src/flowcept/agents/harness/config.py @@ -0,0 +1,176 @@ +"""Runtime configuration for the harness capture path. + +Every knob is an environment variable so that it can be set from a harness +settings file, a plugin ``userConfig``, or a shell profile without needing a +config file on disk. Nothing here imports flowcept. +""" + +from __future__ import annotations + +import os +from pathlib import Path + +ENV_PREFIX = "FLOWCEPT_HARNESS_" + +#: Content capture modes for potentially large tool payloads. +CONTENT_FULL = "full" +CONTENT_SUMMARY = "summary" +CONTENT_NONE = "none" +CONTENT_MODES = (CONTENT_FULL, CONTENT_SUMMARY, CONTENT_NONE) + + +def _env(name: str, default: str | None = None) -> str | None: + return os.environ.get(ENV_PREFIX + name, default) + + +def _env_bool(name: str, default: bool) -> bool: + raw = _env(name) + if raw is None: + return default + return raw.strip().lower() in {"1", "true", "yes", "on"} + + +def _env_int(name: str, default: int) -> int: + raw = _env(name) + if raw is None: + return default + try: + return int(raw) + except ValueError: + return default + + +def default_home() -> Path: + """Return the base directory holding session state and buffers.""" + explicit = _env("HOME") + if explicit: + return Path(explicit).expanduser() + xdg = os.environ.get("XDG_STATE_HOME") + if xdg: + return Path(xdg).expanduser() / "flowcept-harness" + return Path.home() / ".flowcept" / "harness" + + +# Deliberately a plain class rather than a dataclass: importing `dataclasses` +# costs ~80ms (it pulls in `inspect` and `ast`), which is a third of the wall +# time of a hook process that otherwise does almost nothing. +class Config: + """Resolved capture configuration.""" + + __slots__ = ( + "buffer_dir", + "campaign_id", + "campaign_scope", + "capture_prompts", + "capture_telemetry", + "capture_tool_results", + "content_mode", + "debug", + "enabled", + "home", + "max_str", + "online", + "redact", + "source_agent_id", + "timeout_ms", + ) + + def __init__( + self, + enabled: bool = True, + home: Path | None = None, + #: Where records are appended. One JSONL file per session keeps + #: concurrent sessions from interleaving and makes ingest cheap. + buffer_dir: Path | None = None, + #: Group every session under the same project into one campaign. + campaign_id: str | None = None, + campaign_scope: str = "project", # project | global | none + #: Max characters kept for any single captured string. + max_str: int = 4000, + #: How to treat file contents in Write/Edit-style tool inputs. + content_mode: str = CONTENT_SUMMARY, + #: Redact secret-looking keys and obvious credential literals. + redact: bool = True, + capture_telemetry: bool = False, + capture_prompts: bool = True, + capture_tool_results: bool = True, + #: Link every emitted turn/tool/LLM task back to a task or agent from + #: another capture system (stamped as ``source_agent_id``). A + #: ``flowcept_source_agent_id`` hook-payload key wins over this. + source_agent_id: str | None = None, + #: Publish to the Flowcept MQ in addition to the JSONL buffer. + online: bool = False, + #: Hard ceiling on hook wall time. + timeout_ms: int = 2000, + #: Log capture failures instead of staying silent. + debug: bool = False, + ): + self.enabled = enabled + self.home = home if home is not None else default_home() + self.buffer_dir = buffer_dir + self.campaign_id = campaign_id + self.campaign_scope = campaign_scope + self.max_str = max_str + self.content_mode = content_mode + self.redact = redact + self.capture_telemetry = capture_telemetry + self.capture_prompts = capture_prompts + self.capture_tool_results = capture_tool_results + self.source_agent_id = source_agent_id + self.online = online + self.timeout_ms = timeout_ms + self.debug = debug + + def __repr__(self) -> str: + """Return a debug representation listing every configured field.""" + fields = ", ".join(f"{name}={getattr(self, name)!r}" for name in self.__slots__) + return f"Config({fields})" + + @property + def sessions_dir(self) -> Path: + """Directory holding per-session state files.""" + return self.home / "sessions" + + @property + def buffers_dir(self) -> Path: + """Directory holding JSONL buffer files, honoring any override.""" + return self.buffer_dir or (self.home / "buffers") + + @property + def log_path(self) -> Path: + """Path of the harness debug log file.""" + return self.home / "harness.log" + + def buffer_path(self, workflow_id: str) -> Path: + """Return the JSONL buffer file path for ``workflow_id``.""" + return self.buffers_dir / f"{workflow_id}.jsonl" + + def state_path(self, workflow_id: str) -> Path: + """Return the session state file path for ``workflow_id``.""" + return self.sessions_dir / f"{workflow_id}.json" + + +def load_config() -> Config: + """Build a :class:`Config` from the environment.""" + content_mode = (_env("CONTENT", CONTENT_SUMMARY) or CONTENT_SUMMARY).strip().lower() + if content_mode not in CONTENT_MODES: + content_mode = CONTENT_SUMMARY + + buffer_dir = _env("BUFFER_DIR") + return Config( + enabled=_env_bool("ENABLED", True), + home=default_home(), + buffer_dir=Path(buffer_dir).expanduser() if buffer_dir else None, + campaign_id=_env("CAMPAIGN_ID"), + campaign_scope=(_env("CAMPAIGN_SCOPE", "project") or "project").strip().lower(), + max_str=_env_int("MAX_STR", 4000), + content_mode=content_mode, + redact=_env_bool("REDACT", True), + capture_telemetry=_env_bool("TELEMETRY", False), + capture_prompts=_env_bool("CAPTURE_PROMPTS", True), + capture_tool_results=_env_bool("CAPTURE_TOOL_RESULTS", True), + source_agent_id=_env("SOURCE_AGENT_ID"), + online=_env_bool("ONLINE", False), + timeout_ms=_env_int("TIMEOUT_MS", 2000), + debug=_env_bool("DEBUG", False), + ) diff --git a/src/flowcept/agents/harness/emit.py b/src/flowcept/agents/harness/emit.py new file mode 100644 index 00000000..bb030b2e --- /dev/null +++ b/src/flowcept/agents/harness/emit.py @@ -0,0 +1,199 @@ +"""Write provenance records. + +The default (and always-on) sink is a per-session JSONL file whose records are +byte-for-byte the dicts Flowcept itself buffers — ``{"type": "workflow"|"task"| +"agent", ...}``. That means the buffer is directly consumable by +``flowcept --generate-report --input-path `` and by +``flowcept.agents.harness flush`` with no conversion step. + +An optional online sink publishes the same dicts to the Flowcept message queue. +It is off by default because it would add a Redis round-trip to every keystroke- +adjacent hook. +""" + +from __future__ import annotations + +import json +import os +from collections.abc import Iterable +from pathlib import Path +from typing import Any + +from .config import Config +from .state import file_lock + +_HOSTNAME: str | None = None +_SYSTEM: str | None = None +_LOGIN: str | None = None + + +def hostname() -> str: + """Return this machine's name. + + ``os.uname`` is preferred over ``socket.gethostname`` purely for import + cost: ``socket`` is ~11ms to import and this is the only thing it was + needed for. Same for ``getpass`` and ``platform`` below. + """ + global _HOSTNAME + if _HOSTNAME is None: + try: + _HOSTNAME = os.uname().nodename + except AttributeError: # Windows + import socket + + try: + _HOSTNAME = socket.gethostname() + except OSError: + _HOSTNAME = "unknown" + return _HOSTNAME or "unknown" + + +def system_name() -> str: + """Return the OS name, as ``platform.system()`` would spell it.""" + global _SYSTEM + if _SYSTEM is None: + try: + _SYSTEM = os.uname().sysname + except AttributeError: # Windows + _SYSTEM = "Windows" + return _SYSTEM + + +def login_name() -> str: + """Return the current user's login name, cached after the first lookup.""" + global _LOGIN + if _LOGIN is None: + _LOGIN = os.environ.get("USER") or os.environ.get("USERNAME") or "" + if not _LOGIN: + import getpass + + try: + _LOGIN = getpass.getuser() + except Exception: + _LOGIN = "unknown" + return _LOGIN + + +def _default(obj: Any) -> str: + return repr(obj) + + +class JsonlEmitter: + """JSONL sink, safe against concurrent hook processes. + + Appends, except when a record supersedes an earlier one for the same + workflow (see :meth:`write`). + """ + + def __init__(self, path: Path): + self.path = path + + def write(self, records: Iterable[dict[str, Any]], supersede: frozenset[str] | None = None) -> int: + """Write ``records``, dropping superseded workflow records first. + + A workflow is written twice: once when it opens, so an in-flight or + crashed session is still readable, and once when it closes with its + final status and totals. Only the second one should survive. Flowcept's + loader keeps the last workflow record per file, but *only* when the file + holds a single workflow -- with subagents it holds several, and every + record is then counted as a distinct run. So on close the prior record + for that ``workflow_id`` is dropped in a single rewrite rather than + left to be deduplicated downstream. + """ + records = [r for r in records if r] + if not records: + return 0 + self.path.parent.mkdir(parents=True, exist_ok=True) + payload = "".join(json.dumps(r, default=_default, ensure_ascii=False) + "\n" for r in records) + # A single locked append keeps records from interleaving mid-line when + # parallel tool calls flush at the same moment. + with file_lock(self.path): + if supersede and self.path.exists(): + self._rewrite_without(supersede, payload) + else: + with self.path.open("a", encoding="utf-8") as fh: + fh.write(payload) + return len(records) + + def _rewrite_without(self, workflow_ids: frozenset[str], payload: str) -> None: + """Rewrite the buffer without workflow records for ``workflow_ids``. + + Caller must hold the lock. One streaming pass into a temp file, then an + atomic replace, so a reader never sees a partial buffer. Cost is one + file rewrite per session or subagent close, not per event. + """ + tmp = self.path.with_suffix(self.path.suffix + ".tmp") + dropped = False + with self.path.open("r", encoding="utf-8") as src, tmp.open("w", encoding="utf-8") as dst: + for line in src: + stripped = line.strip() + if stripped: + try: + record = json.loads(stripped) + except ValueError: + record = None + if ( + isinstance(record, dict) + and record.get("type") == "workflow" + and record.get("workflow_id") in workflow_ids + ): + dropped = True + continue + dst.write(line) + dst.write(payload) + if dropped: + os.replace(str(tmp), str(self.path)) + else: + # Nothing to drop: the rewrite was wasted work, so keep the cheaper + # append and leave the original file untouched. + tmp.unlink(missing_ok=True) + with self.path.open("a", encoding="utf-8") as fh: + fh.write(payload) + + +class OnlinePublisher: + """Publishes records to the Flowcept MQ. Import of flowcept is deferred.""" + + def __init__(self) -> None: + self._mq = None + + def write(self, records: Iterable[dict[str, Any]]) -> int: + """Publish ``records`` to the MQ, building the connection on first use.""" + records = [r for r in records if r] + if not records: + return 0 + if self._mq is None: + from flowcept.commons.daos.mq_dao.mq_dao_base import MQDao + + self._mq = MQDao.build() + self._mq.bulk_publish(records) + return len(records) + + +class Emitter: + """Fan-out to the configured sinks. Never raises into the harness.""" + + def __init__(self, config: Config, workflow_id: str, on_error=None): + self.config = config + self.jsonl = JsonlEmitter(config.buffer_path(workflow_id)) + self.online = OnlinePublisher() if config.online else None + self._on_error = on_error + + def emit(self, *records: dict[str, Any], supersede: frozenset[str] | None = None) -> int: + """Write ``records`` to every configured sink, swallowing sink errors.""" + flat = [r for r in records if r] + if not flat: + return 0 + written = 0 + try: + written = self.jsonl.write(flat, supersede=supersede) + except Exception as exc: # capture must never break the harness + if self._on_error: + self._on_error(f"jsonl write failed: {exc!r}") + if self.online is not None: + try: + self.online.write(flat) + except Exception as exc: + if self._on_error: + self._on_error(f"online publish failed: {exc!r}") + return written diff --git a/src/flowcept/agents/harness/events.py b/src/flowcept/agents/harness/events.py new file mode 100644 index 00000000..8eafe4df --- /dev/null +++ b/src/flowcept/agents/harness/events.py @@ -0,0 +1,148 @@ +"""The harness-independent event shape. + +Every adapter's only job is to turn its harness's native event into one of +these. The recorder then knows nothing about Claude Code, Codex, or OTel. + +Like :class:`~flowcept.agents.harness.config.Config`, this is a plain slotted class +rather than a dataclass: hooks are short-lived processes and the ``dataclasses`` +import costs more than everything this module does. +""" + +from __future__ import annotations + +import time +from typing import Any + +_FIELDS = ( + "kind", + "harness", + "session_id", + "timestamp", + "started_at", + "cwd", + "project_dir", + "model", + "permission_mode", + "effort", + "source", + "prompt_id", + "prompt", + "response", + "tool_name", + "tool_use_id", + "tool_input", + "tool_response", + "error", + "call_id", + "usage", + "agent_name", + "agent_ref", + "source_agent_id", + "raw", + "message", + "tags", +) + + +class HarnessEvent: + """A normalized lifecycle event from some AI coding harness. + + Attributes + ---------- + kind: + One of :class:`flowcept.agents.harness.vocab.EventKind`. + harness: + Harness identifier, e.g. ``claude_code``, ``codex``, ``langgraph``. + session_id: + The harness's own session identifier. All provenance IDs derive from it. + timestamp: + When the event happened. For an event that completes something (a tool + result, a finished turn) this is the end. + started_at: + When the completed work *began*, for sources that report a duration in + one event rather than a pre/post pair -- an OTel span, or an SDK + callback. ``None`` means the recorder should fall back to the start it + saw earlier, or to ``timestamp``. + source: + Why the event fired: session start reason, end reason, compact trigger. + agent_name: + Subagent type/name, e.g. ``Explore``. ``None`` means the main assistant. + agent_ref: + The harness's own subagent identifier, used to pair start with stop. + source_agent_id: + A task or agent id from *another* capture system (e.g. a + framework-emitted task that launched this session) to link back to. + Hook adapters read it from the ``flowcept_source_agent_id`` payload + key; ``FLOWCEPT_HARNESS_SOURCE_AGENT_ID`` is the env fallback, and the + payload key wins. + raw: + The untouched source event, kept in ``custom_metadata.raw_event``. + """ + + __slots__ = _FIELDS + + def __init__( + self, + kind: str, + harness: str, + session_id: str, + timestamp: float | None = None, + started_at: float | None = None, + cwd: str | None = None, + project_dir: str | None = None, + model: str | None = None, + permission_mode: str | None = None, + effort: str | None = None, + source: str | None = None, + prompt_id: str | None = None, + prompt: str | None = None, + response: str | None = None, + tool_name: str | None = None, + tool_use_id: str | None = None, + tool_input: Any = None, + tool_response: Any = None, + error: str | None = None, + call_id: str | None = None, + usage: dict[str, Any] | None = None, + agent_name: str | None = None, + agent_ref: str | None = None, + source_agent_id: str | None = None, + raw: dict[str, Any] | None = None, + message: str | None = None, + tags: list[str] | None = None, + ): + self.kind = kind + self.harness = harness + self.session_id = session_id + self.timestamp = time.time() if timestamp is None else timestamp + self.started_at = started_at + self.cwd = cwd + self.project_dir = project_dir + self.model = model + self.permission_mode = permission_mode + self.effort = effort + self.source = source + self.prompt_id = prompt_id + self.prompt = prompt + self.response = response + self.tool_name = tool_name + self.tool_use_id = tool_use_id + self.tool_input = tool_input + self.tool_response = tool_response + self.error = error + self.call_id = call_id + self.usage = usage + self.agent_name = agent_name + self.agent_ref = agent_ref + self.source_agent_id = source_agent_id + self.raw = raw + self.message = message + self.tags = tags + + def to_dict(self) -> dict[str, Any]: + """Return the event's non-None fields.""" + return {name: getattr(self, name) for name in _FIELDS if getattr(self, name) is not None} + + def __repr__(self) -> str: + """Return a short debug representation naming the event and session.""" + return f"HarnessEvent(kind={self.kind!r}, harness={self.harness!r}, session_id={self.session_id!r})" diff --git a/src/flowcept/agents/harness/ids.py b/src/flowcept/agents/harness/ids.py new file mode 100644 index 00000000..a4e93423 --- /dev/null +++ b/src/flowcept/agents/harness/ids.py @@ -0,0 +1,76 @@ +"""Deterministic identifier derivation. + +Harness hooks run as independent short-lived processes: the process that sees +``PreToolUse`` is not the process that sees ``PostToolUse``. Rather than pass +IDs through a database, every ID is derived deterministically (UUIDv5) from +identifiers the harness already gives us. Two processes that see the same +``tool_use_id`` therefore compute the same ``task_id`` without coordinating. +""" + +from __future__ import annotations + +import hashlib +import os +import uuid +from pathlib import Path + +#: Stable namespace for all flowcept-harness identifiers. Do not change: it +#: would renumber every previously captured session. +NAMESPACE = uuid.UUID("6f3a3d2e-2b1c-5f4a-9c7e-1a2b3c4d5e6f") + + +def _uuid5(*parts: str) -> str: + return str(uuid.uuid5(NAMESPACE, "|".join(p or "" for p in parts))) + + +def workflow_id_for(harness: str, session_id: str) -> str: + """Workflow ID for one harness session.""" + return _uuid5("workflow", harness, session_id) + + +def agent_id_for(harness: str, session_id: str, agent_name: str | None = None) -> str: + """Agent ID for the assistant driving a session (or a named subagent).""" + return _uuid5("agent", harness, session_id, agent_name or "main") + + +def turn_task_id(workflow_id: str, turn_key: str) -> str: + """Task ID for a prompt->response turn.""" + return _uuid5("turn", workflow_id, turn_key) + + +def tool_task_id(workflow_id: str, tool_key: str) -> str: + """Task ID for a single tool invocation.""" + return _uuid5("tool", workflow_id, tool_key) + + +def llm_task_id(workflow_id: str, call_key: str) -> str: + """Task ID for a single model invocation.""" + return _uuid5("llm", workflow_id, call_key) + + +def subagent_workflow_id(parent_workflow_id: str, agent_key: str) -> str: + """Workflow ID for a subagent, nested under its parent session.""" + return _uuid5("subworkflow", parent_workflow_id, agent_key) + + +def event_task_id(workflow_id: str, kind: str, key: str) -> str: + """Task ID for a point-in-time lifecycle event (compaction, notification).""" + return _uuid5("event", workflow_id, kind, key) + + +def campaign_id_for_project(project_dir: str | os.PathLike[str] | None) -> str: + """Derive a stable campaign ID from a project directory. + + All sessions run against the same checkout land in one campaign, which is + what makes cross-session queries ("every tool call this repo ever caused") + possible without the user having to set anything. + """ + if not project_dir: + return _uuid5("campaign", "default") + resolved = str(Path(project_dir).expanduser().resolve()) + return _uuid5("campaign", resolved) + + +def content_digest(text: str) -> str: + """Short, stable digest used when file contents are summarised, not stored.""" + return hashlib.sha256(text.encode("utf-8", "replace")).hexdigest()[:16] diff --git a/src/flowcept/agents/harness/mcp_server.py b/src/flowcept/agents/harness/mcp_server.py new file mode 100644 index 00000000..e194bce1 --- /dev/null +++ b/src/flowcept/agents/harness/mcp_server.py @@ -0,0 +1,412 @@ +"""An MCP server exposing captured provenance as tools. + +Two audiences: + +*Any MCP client* -- Claude Code, Cursor, an agent of your own -- can ask what +happened in past sessions without knowing the buffer format. That makes +provenance answerable in the same conversation that produced it. + +*Harnesses with no hook system* can call ``record_event`` to push provenance in, +which is the only integration path available when a harness can run an MCP +server but cannot run a command per lifecycle event. + +Run with ``flowcept-harness-mcp``. Requires the ``mcp`` extra; the capture path +does not. +""" + +from __future__ import annotations + +import json +import time +from pathlib import Path +from typing import Any + +from .config import Config, load_config + +SERVER_NAME = "flowcept-provenance" + + +def _load_server_class(): + """Return the MCP server class across SDK generations. + + The SDK renamed ``FastMCP`` to ``MCPServer`` in 2.0 while keeping the + decorator API identical, so supporting both is an import, not a shim. + """ + try: + from mcp.server.mcpserver import MCPServer # SDK >= 2.0 + + return MCPServer + except ImportError: + pass + try: + from mcp.server.fastmcp import FastMCP # SDK < 2.0 + + return FastMCP + except ImportError as exc: # pragma: no cover - depends on the environment + raise SystemExit("The MCP server needs the `mcp` package: pip install 'flowcept-harness[mcp]'") from exc + + +# -- buffer access ----------------------------------------------------------- + + +def _buffers(config: Config) -> list[Path]: + if not config.buffers_dir.is_dir(): + return [] + return sorted(config.buffers_dir.glob("*.jsonl"), key=lambda p: p.stat().st_mtime, reverse=True) + + +def _records(path: Path) -> list[dict[str, Any]]: + out: list[dict[str, Any]] = [] + try: + text = path.read_text(encoding="utf-8") + except OSError: + return out + for line in text.splitlines(): + line = line.strip() + if not line: + continue + try: + record = json.loads(line) + except ValueError: + continue + if isinstance(record, dict): + out.append(record) + return out + + +def _find_buffer(config: Config, session: str | None) -> Path | None: + buffers = _buffers(config) + if not buffers: + return None + if not session: + return buffers[0] + # Accept a full workflow id, a prefix of one, or the harness session id. + for path in buffers: + if path.stem == session or path.stem.startswith(session): + return path + for path in buffers: + for record in _records(path): + if record.get("type") == "workflow" and record.get("used", {}).get("session_id") == session: + return path + return None + + +def _elapsed(record: dict[str, Any]) -> float | None: + started, ended = record.get("started_at"), record.get("ended_at") + if isinstance(started, (int, float)) and isinstance(ended, (int, float)): + return round(ended - started, 3) + return None + + +def build_server(config: Config | None = None): + """Build the MCP server. Separate from :func:`main` so tests can drive it.""" + config = config or load_config() + server_class = _load_server_class() + server = server_class(SERVER_NAME) + + @server.tool( + description=( + "List captured AI coding sessions, newest first. Returns the workflow id, " + "harness, status, timing, and per-session totals for turns, tool calls, " + "tool errors, and subagents." + ) + ) + def list_sessions(limit: int = 20) -> list[dict[str, Any]]: + sessions = [] + for path in _buffers(config)[: max(1, limit)]: + records = _records(path) + root = next( + (r for r in records if r.get("type") == "workflow" and r.get("parent_workflow_id") is None), + None, + ) + if root is None: + continue + sessions.append( + { + "workflow_id": root.get("workflow_id", path.stem), + "harness": (root.get("custom_metadata") or {}).get("harness"), + "name": root.get("name"), + "status": root.get("status"), + "started_at": root.get("started_at"), + "ended_at": root.get("ended_at"), + "elapsed_seconds": _elapsed(root), + "cwd": (root.get("used") or {}).get("cwd"), + "model": (root.get("used") or {}).get("model"), + "totals": root.get("generated") or {}, + } + ) + return sessions + + @server.tool( + description=( + "Get the full activity of one session: every turn, tool call, and subagent " + "in order. `session` accepts a workflow id or prefix; omit it for the most " + "recent session." + ) + ) + def get_session(session: str | None = None, include_io: bool = False) -> dict[str, Any]: + path = _find_buffer(config, session) + if path is None: + return {"error": f"No session matching {session!r}."} + + records = _records(path) + workflows = {r.get("workflow_id"): r for r in records if r.get("type") == "workflow"} + activity = [] + for record in records: + if record.get("type") != "task": + continue + wf = workflows.get(record.get("workflow_id")) or {} + entry = { + "kind": record.get("subtype"), + "name": record.get("activity_id"), + "status": record.get("status"), + "elapsed_seconds": _elapsed(record), + "in_subagent": wf.get("name") if wf.get("parent_workflow_id") else None, + "error": record.get("stderr"), + } + if include_io: + entry["used"] = record.get("used") + entry["generated"] = record.get("generated") + activity.append({k: v for k, v in entry.items() if v is not None}) + + root = next((w for w in workflows.values() if not w.get("parent_workflow_id")), {}) + return { + "workflow_id": root.get("workflow_id", path.stem), + "status": root.get("status"), + "totals": root.get("generated") or {}, + "subagents": [ + {"name": w.get("name"), "status": w.get("status"), "elapsed_seconds": _elapsed(w)} + for w in workflows.values() + if w.get("parent_workflow_id") + ], + "activity": activity, + } + + @server.tool( + description=( + "Search tool calls across captured sessions. Filter by tool name, status " + "('ERROR' finds failures), or a substring of the tool's inputs. Useful for " + "'what commands have I run', 'what has been failing', 'when did I last touch X'." + ) + ) + def search_tool_calls( + tool_name: str | None = None, + status: str | None = None, + contains: str | None = None, + limit: int = 50, + ) -> list[dict[str, Any]]: + hits: list[dict[str, Any]] = [] + for path in _buffers(config): + for record in _records(path): + if record.get("subtype") != "agent_tool": + continue + if tool_name and record.get("activity_id") != tool_name: + continue + if status and record.get("status") != status.upper(): + continue + if contains and contains.lower() not in json.dumps(record.get("used") or {}).lower(): + continue + hits.append( + { + "session": path.stem[:8], + "tool": record.get("activity_id"), + "status": record.get("status"), + "elapsed_seconds": _elapsed(record), + "used": record.get("used"), + "error": record.get("stderr"), + "at": record.get("started_at"), + } + ) + if len(hits) >= limit: + return hits + return hits + + @server.tool( + description=( + "Aggregate statistics over captured sessions: tool call counts and failure " + "rates by tool, total time per tool, and the slowest individual calls." + ) + ) + def session_stats(session: str | None = None) -> dict[str, Any]: + paths = [p for p in ([_find_buffer(config, session)] if session else _buffers(config)) if p] + by_tool: dict[str, dict[str, Any]] = {} + slowest: list[dict[str, Any]] = [] + + for path in paths: + for record in _records(path): + if record.get("subtype") != "agent_tool": + continue + name = record.get("activity_id", "?") + stats = by_tool.setdefault(name, {"calls": 0, "errors": 0, "seconds": 0.0}) + stats["calls"] += 1 + if record.get("status") == "ERROR": + stats["errors"] += 1 + elapsed = _elapsed(record) + if elapsed is not None: + stats["seconds"] = round(stats["seconds"] + elapsed, 3) + slowest.append({"tool": name, "seconds": elapsed, "session": path.stem[:8]}) + + slowest.sort(key=lambda r: r["seconds"], reverse=True) + return { + "sessions": len(paths), + "by_tool": dict(sorted(by_tool.items(), key=lambda kv: kv[1]["calls"], reverse=True)), + "slowest_calls": slowest[:10], + } + + # -- provenance analysis tools (thin wrappers over prov_analysis.core) ---- + + def _session_records(session: str | None): + path = _find_buffer(config, session) + if path is None: + return None, {"error": f"No session matching {session!r}."} + return _records(path), None + + @server.tool( + description=( + "Analyze one captured session end to end: execution summary (counts by " + "activity/subtype, statuses, duration bounds, token usage) plus per-agent " + "behavior (turns, tool calls by tool, LLM calls, subagents). `session` " + "accepts a workflow id or prefix; omit it for the most recent session." + ) + ) + def analyze_session(session: str | None = None) -> dict[str, Any]: + from flowcept.agents.prov_analysis import core as prov_core + + records, error = _session_records(session) + if error: + return error + return { + "summary": prov_core.summarize_execution(records), + "agent_behavior": prov_core.analyze_agent_behavior(records), + } + + @server.tool( + description=( + "Analyze failures in one captured session: failed tasks grouped by " + "activity with error excerpts, error rate per activity, and first/last " + "failure times. Omit `session` for the most recent session." + ) + ) + def analyze_errors(session: str | None = None) -> dict[str, Any]: + from flowcept.agents.prov_analysis import core as prov_core + + records, error = _session_records(session) + if error: + return error + return prov_core.analyze_errors(records) + + @server.tool( + description=( + "Find the slowest tasks of one captured session, longest elapsed first, " + "with activity, status, and parent-chain depth. Omit `session` for the " + "most recent session." + ) + ) + def find_slowest(session: str | None = None, limit: int = 10) -> list[dict[str, Any]]: + from flowcept.agents.prov_analysis import core as prov_core + + records, error = _session_records(session) + if error: + return [error] + return prov_core.find_slowest_tasks(records, limit=limit) + + @server.tool( + description=( + "List cross-framework provenance links in one captured session: edges " + "built from source_agent_id pointers between frameworks (e.g. a LangGraph " + "run linked to a task from another agent framework), plus the count of " + "unlinked tasks. Omit `session` for the most recent session." + ) + ) + def cross_links(session: str | None = None) -> dict[str, Any]: + from flowcept.agents.prov_analysis import core as prov_core + + records, error = _session_records(session) + if error: + return error + return prov_core.cross_framework_links(records) + + @server.tool( + description=( + "Record a provenance event from a harness that has no hook system. `kind` is " + "one of session_start, prompt, tool_pre, tool_post, tool_error, llm_call, " + "turn_end, subagent_start, subagent_stop, session_end. Pass the same " + "`session_id` and `harness` for every event in one session: together they " + "identify the session, so changing either starts a separate one." + ) + ) + def record_event( + kind: str, + session_id: str, + harness: str = "mcp", + tool_name: str | None = None, + tool_input: dict[str, Any] | None = None, + tool_response: dict[str, Any] | None = None, + prompt: str | None = None, + response: str | None = None, + model: str | None = None, + error: str | None = None, + ) -> dict[str, Any]: + from .events import HarnessEvent + from .recorder import Recorder + from .vocab import EventKind + + if kind not in EventKind.ALL: + return {"error": f"Unknown kind {kind!r}. Valid kinds: {sorted(EventKind.ALL)}"} + + records = Recorder(config).record( + HarnessEvent( + kind=kind, + harness=harness, + session_id=session_id, + timestamp=time.time(), + model=model, + prompt=prompt, + response=response, + tool_name=tool_name, + tool_input=tool_input, + tool_response=tool_response, + error=error, + ) + ) + return {"recorded": len(records)} + + @server.tool( + description=( + "Generate a Flowcept workflow card for a session: a markdown report with " + "timings, status counts, slowest activities, and per-workflow detail." + ) + ) + def generate_report(session: str | None = None) -> dict[str, Any]: + path = _find_buffer(config, session) + if path is None: + return {"error": f"No session matching {session!r}."} + try: + from flowcept import Flowcept + except ImportError: + return {"error": "flowcept is not installed; pip install 'flowcept-harness[ingest]'"} + + stats = Flowcept.generate_report( + report_type="workflow_card", + input_jsonl_path=str(path), + format="markdown", + ) + return {"markdown": stats.get("markdown"), "workflows": stats.get("n_workflows"), "tasks": stats.get("n_tasks")} + + return server + + +def main(argv: list[str] | None = None) -> int: + """Entry point for ``flowcept-harness-mcp``.""" + import argparse + + parser = argparse.ArgumentParser(prog="flowcept-harness-mcp", description=__doc__) + parser.add_argument("--transport", default="stdio", choices=("stdio", "sse", "streamable-http")) + args = parser.parse_args(argv) + + build_server().run(transport=args.transport) + return 0 + + +if __name__ == "__main__": # pragma: no cover + raise SystemExit(main()) diff --git a/src/flowcept/agents/harness/prov.py b/src/flowcept/agents/harness/prov.py new file mode 100644 index 00000000..30d2a45b --- /dev/null +++ b/src/flowcept/agents/harness/prov.py @@ -0,0 +1,154 @@ +"""Builders for Flowcept-shaped provenance records. + +These produce exactly the dicts Flowcept buffers internally, so they can be +inserted by the document inserter, published to the MQ, or read back by +``flowcept --generate-report`` without translation. Keys with ``None`` values +are dropped, matching ``TaskObject.to_dict``. +""" + +from __future__ import annotations + +import sys +import time +from typing import Any + +from . import vocab +from .emit import hostname, login_name, system_name + + +def _clean(record: dict[str, Any]) -> dict[str, Any]: + return {k: v for k, v in record.items() if v is not None} + + +def _env_fields() -> dict[str, Any]: + return { + "hostname": hostname(), + "node_name": hostname(), + "login_name": login_name(), + } + + +def workflow_record( + *, + workflow_id: str, + campaign_id: str | None = None, + name: str | None = None, + subtype: str = vocab.AGENT_SESSION, + agent_id: str | None = None, + parent_workflow_id: str | None = None, + used: dict[str, Any] | None = None, + generated: dict[str, Any] | None = None, + custom_metadata: dict[str, Any] | None = None, + started_at: float | None = None, + ended_at: float | None = None, + status: str | None = None, + description: str | None = None, +) -> dict[str, Any]: + """Build a ``type: workflow`` record. + + Flowcept upserts workflows by ``workflow_id``, and the JSONL report loader + keeps the last record for an id, so emitting an updated record at session + end is the supported way to close a workflow out. + """ + return _clean( + { + "type": vocab.TYPE_WORKFLOW, + "workflow_id": workflow_id, + "parent_workflow_id": parent_workflow_id, + "campaign_id": campaign_id, + "name": name, + "subtype": subtype, + "agent_id": agent_id, + "adapter_id": vocab.ADAPTER_ID, + "user": login_name(), + "utc_timestamp": time.time(), + "started_at": started_at, + "ended_at": ended_at, + "status": status, + "workflow_description": description, + "used": used, + "generated": generated, + "custom_metadata": custom_metadata, + "environment_id": f"python{sys.version_info.major}.{sys.version_info.minor}-{system_name().lower()}", + "sys_name": system_name(), + } + ) + + +def agent_record( + *, + agent_id: str, + name: str, + workflow_id: str | None = None, + campaign_id: str | None = None, + extra_metadata: dict[str, Any] | None = None, +) -> dict[str, Any]: + """Build a ``type: agent`` record describing the assistant behind a session.""" + return _clean( + { + "type": vocab.TYPE_AGENT, + "agent_id": agent_id, + "name": name, + "workflow_id": workflow_id, + "campaign_id": campaign_id, + "user": login_name(), + "registered_at": time.time(), + "extra_metadata": extra_metadata, + } + ) + + +def task_record( + *, + task_id: str, + workflow_id: str, + activity_id: str, + subtype: str, + campaign_id: str | None = None, + agent_id: str | None = None, + source_agent_id: str | None = None, + parent_task_id: str | None = None, + used: dict[str, Any] | None = None, + generated: dict[str, Any] | None = None, + custom_metadata: dict[str, Any] | None = None, + status: str = vocab.STATUS_FINISHED, + started_at: float | None = None, + ended_at: float | None = None, + stdout: Any = None, + stderr: Any = None, + tags: list[str] | None = None, + dependencies: list[str] | None = None, + telemetry_at_start: dict[str, Any] | None = None, + telemetry_at_end: dict[str, Any] | None = None, +) -> dict[str, Any]: + """Build a ``type: task`` record.""" + now = time.time() + return _clean( + { + "type": vocab.TYPE_TASK, + "task_id": task_id, + "workflow_id": workflow_id, + "campaign_id": campaign_id, + "activity_id": activity_id, + "subtype": subtype, + "adapter_id": vocab.ADAPTER_ID, + "agent_id": agent_id, + "source_agent_id": source_agent_id, + "parent_task_id": parent_task_id, + "used": used, + "generated": generated, + "custom_metadata": custom_metadata, + "status": status, + "started_at": started_at if started_at is not None else now, + "ended_at": ended_at if ended_at is not None else now, + "utc_timestamp": now, + "stdout": stdout, + "stderr": stderr, + "tags": tags, + "dependencies": dependencies, + "telemetry_at_start": telemetry_at_start, + "telemetry_at_end": telemetry_at_end, + "user": login_name(), + **_env_fields(), + } + ) diff --git a/src/flowcept/agents/harness/recorder.py b/src/flowcept/agents/harness/recorder.py new file mode 100644 index 00000000..11017ecc --- /dev/null +++ b/src/flowcept/agents/harness/recorder.py @@ -0,0 +1,604 @@ +"""The state machine that turns harness events into PROV-AGENT records. + +This is the only place that knows how a coding session decomposes into +provenance: + + session -> workflow (subtype ``agent_session``) + turn -> task (subtype ``ai_model_invocation``, granularity=turn) + tool call -> task (subtype ``agent_tool``, parent = the turn) + subagent -> nested workflow + task + lifecycle -> task (subtype ``harness_event``) + +The parent/child edges are what make the capture useful: given a bad file +edit you can walk up to the turn that caused it and the prompt that started +it, which is the ``wasInformedBy`` chain PROV-AGENT is built around. +""" + +from __future__ import annotations + +import time +from typing import Any + +from . import ids, prov, vocab +from .config import Config, load_config +from .emit import Emitter +from .events import HarnessEvent +from .sanitize import Sanitizer +from .state import session_state +from .vocab import EventKind + + +class Recorder: + """Applies one :class:`HarnessEvent` to a session's provenance.""" + + def __init__(self, config: Config | None = None, on_error=None): + self.config = config or load_config() + self.sanitizer = Sanitizer(self.config) + self._on_error = on_error + self._supersede: set[str] = set() + + # -- public API ---------------------------------------------------------- + + def record(self, event: HarnessEvent) -> list[dict[str, Any]]: + """Process an event; return the records that were emitted.""" + if not self.config.enabled: + return [] + + workflow_id = ids.workflow_id_for(event.harness, event.session_id) + emitter = Emitter(self.config, workflow_id, on_error=self._on_error) + + # Workflow ids whose "open" record this event's records replace; see + # JsonlEmitter.write. + self._supersede: set[str] = set() + + with session_state(self.config.state_path(workflow_id)) as state: + records = self._dispatch(event, workflow_id, state) + + if records: + emitter.emit(*records, supersede=frozenset(self._supersede) or None) + return records + + # -- dispatch ------------------------------------------------------------ + + def _dispatch(self, event: HarnessEvent, workflow_id: str, state) -> list[dict[str, Any]]: + records: list[dict[str, Any]] = [] + + # Every event can be the first one we see: a hook may be installed + # mid-session, or SessionStart may not exist on this harness. + records.extend(self._ensure_session(event, workflow_id, state)) + + handler = { + EventKind.SESSION_START: self._noop, + EventKind.PROMPT: self._on_prompt, + EventKind.TURN_END: self._on_turn_end, + EventKind.TOOL_PRE: self._on_tool_pre, + EventKind.TOOL_POST: self._on_tool_post, + EventKind.TOOL_ERROR: self._on_tool_post, + EventKind.LLM_CALL: self._on_llm_call, + EventKind.SUBAGENT_START: self._on_subagent_start, + EventKind.SUBAGENT_STOP: self._on_subagent_stop, + EventKind.NOTIFICATION: self._on_lifecycle_event, + EventKind.COMPACT: self._on_lifecycle_event, + EventKind.SESSION_END: self._on_session_end, + }.get(event.kind) + + if handler is None: + self._warn(f"unknown event kind {event.kind!r}") + return records + + records.extend(handler(event, workflow_id, state)) + return records + + # -- session ------------------------------------------------------------- + + def _ensure_session(self, event: HarnessEvent, workflow_id: str, state) -> list[dict[str, Any]]: + campaign_id = self._campaign_id(event, state) + agent_id = ids.agent_id_for(event.harness, event.session_id) + + # Late-arriving context (the model is only on SessionStart) is folded in + # as it becomes known. + for key, value in ( + ("harness", event.harness), + ("session_id", event.session_id), + ("workflow_id", workflow_id), + ("campaign_id", campaign_id), + ("agent_id", agent_id), + # Only the harness's own cwd, never the fallback derived from the + # hook process, which may run from somewhere else entirely. + ("cwd", event.cwd), + ("project_dir", event.project_dir), + ("model", event.model), + ): + if value is not None: + state.set(key, value) + state.setdefault("started_at", event.timestamp) + state.setdefault("counters", {}) + if event.kind == EventKind.SESSION_START and event.source: + state.setdefault("start_reason", event.source) + + if not state.mark_once("session_opened"): + return [] + + name = f"{event.harness} session" + used = self._clean_dict( + { + "cwd": event.cwd, + "project_dir": event.project_dir, + "model": event.model, + "session_id": event.session_id, + "start_reason": state.get("start_reason"), + } + ) + return [ + prov.agent_record( + agent_id=agent_id, + name=f"{event.harness}:{event.model or 'assistant'}", + workflow_id=workflow_id, + campaign_id=campaign_id, + extra_metadata={"harness": event.harness, "session_id": event.session_id}, + ), + prov.workflow_record( + workflow_id=workflow_id, + campaign_id=campaign_id, + name=name, + subtype=vocab.AGENT_SESSION, + agent_id=agent_id, + used=used, + custom_metadata={"harness": event.harness, "capture": "flowcept-harness"}, + started_at=event.timestamp, + status=vocab.STATUS_RUNNING, + description=f"Interactive {event.harness} session captured by flowcept-harness.", + ), + ] + + def _on_session_end(self, event: HarnessEvent, workflow_id: str, state) -> list[dict[str, Any]]: + records = self._close_open_turn(event, workflow_id, state, reason="session_ended") + counters = state.get("counters", {}) + # This record replaces the one written at session start, so it has to + # carry that record's fields too. + self._supersede.add(workflow_id) + records.append( + prov.workflow_record( + workflow_id=workflow_id, + campaign_id=state.get("campaign_id"), + name=f"{event.harness} session", + subtype=vocab.AGENT_SESSION, + agent_id=state.get("agent_id"), + used=self._clean_dict( + { + "cwd": state.get("cwd"), + "project_dir": state.get("project_dir"), + "model": state.get("model"), + "session_id": event.session_id, + "start_reason": state.get("start_reason"), + } + ), + generated=self._clean_dict( + { + "turns": counters.get("turns"), + "tool_calls": counters.get("tools"), + "tool_errors": counters.get("tool_errors"), + "subagents": counters.get("subagents"), + } + ), + custom_metadata={"harness": event.harness, "end_reason": event.source}, + started_at=state.get("started_at"), + ended_at=event.timestamp, + status=vocab.STATUS_FINISHED, + ) + ) + state.set("ended_at", event.timestamp) + return records + + # -- turns --------------------------------------------------------------- + + def _on_prompt(self, event: HarnessEvent, workflow_id: str, state): + # A prompt while a turn is open means the previous turn never closed. + records = self._close_open_turn(event, workflow_id, state, reason="superseded") + + turn_number = int(state.get("counters", {}).get("turns", 0)) + 1 + turn_key = event.prompt_id or f"n{turn_number}" + task_id = ids.turn_task_id(workflow_id, turn_key) + + state.set( + "current_turn", + { + "task_id": task_id, + "turn_key": turn_key, + "number": turn_number, + "started_at": event.timestamp, + "prompt": self._prompt_value(event.prompt), + "tool_task_ids": [], + "permission_mode": event.permission_mode, + "effort": event.effort, + }, + ) + self._bump(state, "turns") + return records + + def _on_turn_end(self, event: HarnessEvent, workflow_id: str, state): + turn = state.get("current_turn") + if not turn: + # Stop without a matching prompt (hook installed mid-turn). Record + # the response alone rather than dropping it. + turn = { + "task_id": ids.turn_task_id(workflow_id, f"orphan-{int(event.timestamp * 1000)}"), + "number": None, + "started_at": event.timestamp, + "prompt": None, + "tool_task_ids": [], + } + state.set("current_turn", None) + status = vocab.STATUS_ERROR if event.error else vocab.STATUS_FINISHED + return [self._turn_task(event, workflow_id, state, turn, event.response, status)] + + def _close_open_turn(self, event: HarnessEvent, workflow_id: str, state, *, reason: str): + turn = state.get("current_turn") + if not turn: + return [] + state.set("current_turn", None) + return [ + self._turn_task( + event, + workflow_id, + state, + turn, + None, + vocab.STATUS_UNKNOWN, + # The turn is closed, just not by a turn-end event: the user + # interrupted, or the session went away underneath it. + extra_metadata={"close_reason": reason}, + ) + ] + + def _turn_task(self, event, workflow_id, state, turn, response, status, extra_metadata=None): + metadata = { + "granularity": "turn", + "harness": event.harness, + "turn_number": turn.get("number"), + "model": state.get("model"), + "permission_mode": turn.get("permission_mode"), + "effort": turn.get("effort"), + "tool_task_ids": turn.get("tool_task_ids") or None, + "tool_call_count": len(turn.get("tool_task_ids") or []), + } + if extra_metadata: + metadata.update(extra_metadata) + if event.usage: + metadata["llm_usage"] = self.sanitizer.mapping(event.usage) + + generated = None + if response is not None: + generated = {"response": self._prompt_value(response)} + + return prov.task_record( + task_id=turn["task_id"], + workflow_id=workflow_id, + campaign_id=state.get("campaign_id"), + activity_id="agent_turn", + subtype=vocab.AI_MODEL_INVOCATION, + agent_id=state.get("agent_id"), + source_agent_id=self._source_agent_id(event), + used=self._clean_dict({"prompt": turn.get("prompt")}), + generated=generated, + custom_metadata=self._clean_dict(metadata), + status=status, + started_at=turn.get("started_at"), + ended_at=event.timestamp, + stderr=event.error, + ) + + # -- tools --------------------------------------------------------------- + + def _tool_key(self, event: HarnessEvent) -> str: + return event.tool_use_id or f"{event.tool_name}@{event.timestamp}" + + def _on_tool_pre(self, event: HarnessEvent, workflow_id: str, state): + turn = state.get("current_turn") or {} + state.add_pending( + self._tool_key(event), + { + "started_at": event.timestamp, + "tool_name": event.tool_name, + "used": self.sanitizer.mapping(event.tool_input) if event.tool_input is not None else None, + "parent_task_id": turn.get("task_id"), + "agent_ref": event.agent_ref, + }, + ) + return [] + + def _on_tool_post(self, event: HarnessEvent, workflow_id: str, state): + key = self._tool_key(event) + pending = state.pop_pending(key) or {} + task_id = ids.tool_task_id(workflow_id, key) + + turn = state.get("current_turn") or {} + parent_task_id = pending.get("parent_task_id") or turn.get("task_id") + tool_name = event.tool_name or pending.get("tool_name") or "unknown_tool" + + used = pending.get("used") + if used is None and event.tool_input is not None: + used = self.sanitizer.mapping(event.tool_input) + + generated = None + if self.config.capture_tool_results and event.tool_response is not None: + generated = self.sanitizer.mapping(event.tool_response) + + is_error = event.kind == EventKind.TOOL_ERROR or bool(event.error) + status = vocab.STATUS_ERROR if is_error else vocab.STATUS_FINISHED + + # Record the tool on its turn so the turn can list what it caused. + if turn and turn.get("task_id") == parent_task_id: + turn.setdefault("tool_task_ids", []).append(task_id) + state.set("current_turn", turn) + + self._bump(state, "tools") + if is_error: + self._bump(state, "tool_errors") + + # A subagent's tool call belongs to the subagent's workflow, not the + # session's, so the two do not interleave in dataflow views. + target_workflow_id = workflow_id + agent_id = state.get("agent_id") + if event.agent_ref: + sub = (state.get("subagents") or {}).get(event.agent_ref) + if sub: + target_workflow_id = sub.get("workflow_id", workflow_id) + agent_id = sub.get("agent_id", agent_id) + parent_task_id = sub.get("task_id") or parent_task_id + + return [ + prov.task_record( + task_id=task_id, + workflow_id=target_workflow_id, + campaign_id=state.get("campaign_id"), + activity_id=tool_name, + subtype=vocab.AGENT_TOOL, + agent_id=agent_id, + source_agent_id=self._source_agent_id(event), + parent_task_id=parent_task_id, + used=used, + generated=generated, + custom_metadata=self._clean_dict( + { + "harness": event.harness, + "tool_name": tool_name, + "tool_use_id": event.tool_use_id, + "permission_mode": event.permission_mode, + "mcp_tool": bool(tool_name.startswith("mcp__")), + "duration_known": "started_at" in pending, + } + ), + status=status, + # An event that reports its own duration (an OTel span) wins + # over the start recorded by a matching pre-event, which wins + # over "we only ever saw the end". + started_at=event.started_at or pending.get("started_at", event.timestamp), + ended_at=event.timestamp, + stderr=event.error, + tags=event.tags, + ) + ] + + # -- model invocations --------------------------------------------------- + + def _on_llm_call(self, event: HarnessEvent, workflow_id: str, state): + turn = state.get("current_turn") or {} + call_key = event.call_id or f"{event.timestamp}" + usage = self.sanitizer.mapping(event.usage) if event.usage else None + metadata = self._clean_dict( + { + "granularity": "call", + "harness": event.harness, + "model": event.model or state.get("model"), + "llm_usage": usage, + "provider_request_id": event.call_id, + } + ) + self._bump(state, "llm_calls") + return [ + prov.task_record( + task_id=ids.llm_task_id(workflow_id, call_key), + workflow_id=workflow_id, + campaign_id=state.get("campaign_id"), + activity_id="llm_interaction", + subtype=vocab.AI_MODEL_INVOCATION, + agent_id=state.get("agent_id"), + source_agent_id=self._source_agent_id(event), + parent_task_id=turn.get("task_id"), + used=self._clean_dict({"prompt": self._prompt_value(event.prompt)}), + generated=self._clean_dict({"response": self._prompt_value(event.response)}), + custom_metadata=metadata, + status=vocab.STATUS_ERROR if event.error else vocab.STATUS_FINISHED, + started_at=event.started_at or event.timestamp, + ended_at=event.timestamp, + stderr=event.error, + ) + ] + + # -- subagents ----------------------------------------------------------- + + def _on_subagent_start(self, event: HarnessEvent, workflow_id: str, state): + ref = event.agent_ref or event.agent_name or f"sub-{event.timestamp}" + sub_workflow_id = ids.subagent_workflow_id(workflow_id, ref) + sub_agent_id = ids.agent_id_for(event.harness, event.session_id, event.agent_name or ref) + turn = state.get("current_turn") or {} + + subagents = state.get("subagents") or {} + subagents[ref] = { + "workflow_id": sub_workflow_id, + "agent_id": sub_agent_id, + "agent_name": event.agent_name, + "started_at": event.timestamp, + "task_id": turn.get("task_id"), + # Kept so the closing record can restate what the open one said. + "prompt": self._prompt_value(event.prompt), + } + state.set("subagents", subagents) + self._bump(state, "subagents") + + return [ + prov.agent_record( + agent_id=sub_agent_id, + name=f"{event.harness}:{event.agent_name or 'subagent'}", + workflow_id=sub_workflow_id, + campaign_id=state.get("campaign_id"), + extra_metadata={"harness": event.harness, "subagent_of": state.get("agent_id")}, + ), + prov.workflow_record( + workflow_id=sub_workflow_id, + parent_workflow_id=workflow_id, + campaign_id=state.get("campaign_id"), + name=f"subagent:{event.agent_name or ref}", + subtype=vocab.SUBAGENT_SESSION, + agent_id=sub_agent_id, + used=self._clean_dict({"agent_type": event.agent_name, "prompt": self._prompt_value(event.prompt)}), + custom_metadata=self._clean_dict({"harness": event.harness, "spawned_by_task_id": turn.get("task_id")}), + started_at=event.timestamp, + status=vocab.STATUS_RUNNING, + ), + ] + + def _on_subagent_stop(self, event: HarnessEvent, workflow_id: str, state): + ref = event.agent_ref or event.agent_name or "" + subagents = state.get("subagents") or {} + sub = subagents.pop(ref, None) + state.set("subagents", subagents) + + if sub is None: + sub = { + "workflow_id": ids.subagent_workflow_id(workflow_id, ref or f"sub-{event.timestamp}"), + "agent_id": ids.agent_id_for(event.harness, event.session_id, event.agent_name or ref), + "started_at": event.timestamp, + } + + agent_name = event.agent_name or sub.get("agent_name") + self._supersede.add(sub["workflow_id"]) + return [ + prov.workflow_record( + workflow_id=sub["workflow_id"], + parent_workflow_id=workflow_id, + campaign_id=state.get("campaign_id"), + name=f"subagent:{agent_name or ref}", + subtype=vocab.SUBAGENT_SESSION, + agent_id=sub.get("agent_id"), + used=self._clean_dict({"agent_type": agent_name, "prompt": sub.get("prompt")}), + generated=self._clean_dict({"response": self._prompt_value(event.response)}), + custom_metadata=self._clean_dict( + { + "harness": event.harness, + "agent_type": agent_name, + "spawned_by_task_id": sub.get("task_id"), + } + ), + started_at=sub.get("started_at"), + ended_at=event.timestamp, + status=vocab.STATUS_ERROR if event.error else vocab.STATUS_FINISHED, + ) + ] + + # -- lifecycle ----------------------------------------------------------- + + def _on_lifecycle_event(self, event: HarnessEvent, workflow_id: str, state): + turn = state.get("current_turn") or {} + key = f"{event.kind}-{event.timestamp}" + return [ + prov.task_record( + task_id=ids.event_task_id(workflow_id, event.kind, key), + workflow_id=workflow_id, + campaign_id=state.get("campaign_id"), + activity_id=event.kind, + subtype=vocab.HARNESS_EVENT, + agent_id=state.get("agent_id"), + parent_task_id=turn.get("task_id"), + used=self._clean_dict({"trigger": event.source, "message": event.message}), + custom_metadata=self._clean_dict({"harness": event.harness, "raw_event": self._raw(event)}), + status=vocab.STATUS_FINISHED, + started_at=event.timestamp, + ended_at=event.timestamp, + ) + ] + + def _noop(self, event, workflow_id, state): + return [] + + # -- helpers ------------------------------------------------------------- + + def _campaign_id(self, event: HarnessEvent, state) -> str | None: + existing = state.get("campaign_id") + if existing: + return existing + if self.config.campaign_id: + return self.config.campaign_id + scope = self.config.campaign_scope + if scope == "none": + return None + if scope == "global": + return ids.campaign_id_for_project(None) + return ids.campaign_id_for_project(event.project_dir or event.cwd) + + def _source_agent_id(self, event: HarnessEvent) -> str | None: + """Cross-system source id for this event's task records. + + The ``flowcept_source_agent_id`` payload key (carried on the event) + wins over the ``FLOWCEPT_HARNESS_SOURCE_AGENT_ID`` environment value + (carried on the config). Only turn, tool, and LLM-call tasks get it: + lifecycle events have no dataflow of their own, and session/subagent + workflow records do not model the field. + """ + return event.source_agent_id or self.config.source_agent_id + + def _prompt_value(self, text: str | None): + if text is None: + return None + if not self.config.capture_prompts: + return {"_summary": True, "chars": len(text), "sha256_16": ids.content_digest(text)} + return self.sanitizer.value(text) + + def _raw(self, event: HarnessEvent): + if not event.raw: + return None + return self.sanitizer.value(event.raw) + + @staticmethod + def _clean_dict(data: dict[str, Any]) -> dict[str, Any] | None: + cleaned = {k: v for k, v in data.items() if v is not None} + return cleaned or None + + @staticmethod + def _bump(state, counter: str, amount: int = 1) -> None: + counters = state.get("counters", {}) + counters[counter] = int(counters.get(counter, 0)) + amount + state.set("counters", counters) + + def _warn(self, message: str) -> None: + if self._on_error: + self._on_error(message) + + +def repair_session(config: Config, workflow_id: str) -> list[dict[str, Any]]: + """Close a session whose harness exited without a session-end event. + + Emits the dangling turn (if any) and a terminal workflow record so the + buffer is complete even after a crash or a ``kill -9``. + """ + state_path = config.state_path(workflow_id) + if not state_path.exists(): + return [] + + recorder = Recorder(config) + with session_state(state_path) as state: + if state.get("ended_at"): + return [] + event = HarnessEvent( + kind=EventKind.SESSION_END, + harness=state.get("harness") or "unknown", + session_id=state.get("session_id") or workflow_id, + timestamp=time.time(), + source="repaired", + ) + records = recorder._on_session_end(event, workflow_id, state) + + if records: + Emitter(config, workflow_id).emit(*records, supersede=frozenset(recorder._supersede) or None) + return records diff --git a/src/flowcept/agents/harness/runtime.py b/src/flowcept/agents/harness/runtime.py new file mode 100644 index 00000000..5f79d5e3 --- /dev/null +++ b/src/flowcept/agents/harness/runtime.py @@ -0,0 +1,144 @@ +"""Guard rails for code that runs on a harness's critical path. + +A capture hook has one hard obligation: never make the harness worse. It must +not block the UI, must not write to stdout (some events feed stdout straight +into the model's context), and must not fail in a way the harness reports as an +error. Everything here exists to enforce that. +""" + +from __future__ import annotations + +import os +import signal +import sys +import threading +import time +import traceback +from collections.abc import Callable +from pathlib import Path +from typing import Any + +from .config import Config, load_config + + +def log_error(config: Config, message: str) -> None: + """Append a capture failure to the harness log. Never raises.""" + try: + config.home.mkdir(parents=True, exist_ok=True) + with config.log_path.open("a", encoding="utf-8") as fh: + fh.write(f"{time.strftime('%Y-%m-%dT%H:%M:%S')} pid={os.getpid()} {message}\n") + except Exception: + pass + + +class _Watchdog: + """Hard-exit the process if capture overruns its budget. + + A hook that hangs (a stuck lock, a wedged network sink) stalls the harness. + Losing one provenance record is strictly better than that, so past the + deadline we abandon the work and exit cleanly. + """ + + def __init__(self, timeout_ms: int): + self.timeout = max(timeout_ms, 100) / 1000.0 + self._timer: threading.Timer | None = None + + def __enter__(self): + if self.timeout <= 0: + return self + self._timer = threading.Timer(self.timeout, self._fire) + self._timer.daemon = True + self._timer.start() + return self + + def __exit__(self, *exc): + if self._timer: + self._timer.cancel() + return False + + @staticmethod + def _fire(): + # os._exit skips atexit/flush handlers on purpose: we are already past + # the point where an orderly shutdown is affordable. + os._exit(0) + + +def run_capture(fn: Callable[[Config], Any], config: Config | None = None) -> int: + """Run ``fn`` under the capture safety contract; always return 0.""" + try: + config = config or load_config() + except Exception: + return 0 + + if not config.enabled: + return 0 + + try: + with _Watchdog(config.timeout_ms): + fn(config) + except Exception: + log_error(config, "capture failed:\n" + traceback.format_exc()) + if config.debug: + # Stderr on exit 0 goes to the harness debug log only, never to the + # model, so this is safe to surface when explicitly debugging. + print(traceback.format_exc(), file=sys.stderr) + return 0 + + +def read_stdin_json() -> dict[str, Any]: + """Read a JSON object from stdin; return ``{}`` on anything unexpected.""" + import json + + try: + raw = sys.stdin.read() + except Exception: + return {} + raw = (raw or "").strip() + if not raw: + return {} + try: + data = json.loads(raw) + except json.JSONDecodeError: + return {} + return data if isinstance(data, dict) else {} + + +def detach() -> bool: + """Fork a child to finish the work and return control to the harness now. + + Only used when the sink may be slow (online publishing). Returns True in + the parent (which should return immediately) and False in the child. + """ + if not hasattr(os, "fork"): + return False + try: + pid = os.fork() + except OSError: + return False + if pid > 0: + return True + # Child: detach from the harness's process group so it is not killed with it. + try: + os.setsid() + signal.signal(signal.SIGHUP, signal.SIG_IGN) + devnull = os.open(os.devnull, os.O_RDWR) + os.dup2(devnull, 0) + os.dup2(devnull, 1) + os.dup2(devnull, 2) + except Exception: + pass + return False + + +def project_dir_from_env(fallback: str | None = None) -> str | None: + """Best-effort project root, preferring what the harness told us.""" + for var in ("CLAUDE_PROJECT_DIR", "FLOWCEPT_HARNESS_PROJECT_DIR"): + value = os.environ.get(var) + if value: + return value + if fallback: + return fallback + try: + return str(Path.cwd()) + except OSError: + return None diff --git a/src/flowcept/agents/harness/sanitize.py b/src/flowcept/agents/harness/sanitize.py new file mode 100644 index 00000000..3df7dafc --- /dev/null +++ b/src/flowcept/agents/harness/sanitize.py @@ -0,0 +1,164 @@ +"""Make harness payloads safe and small enough to store as provenance. + +Three concerns, in order of importance: + +1. **Secrets.** Tool inputs routinely contain API keys (a ``Bash`` command with + an inline token, an ``env`` dict, a ``.env`` file write). Provenance is + long-lived and often shared, so redaction happens at capture time, not at + query time. +2. **Size.** A single ``Write`` can carry a megabyte. Provenance wants the + shape of the dataflow, not a second copy of the repository. +3. **JSON-safety.** Whatever we emit must survive ``json.dumps``. +""" + +from __future__ import annotations + +import json +import re +from typing import Any + +from .config import CONTENT_FULL, CONTENT_NONE, CONTENT_SUMMARY, Config +from .ids import content_digest + +REDACTED = "«redacted»" + +#: Key names whose *values* are always dropped. +#: +#: ``auth`` is anchored so it does not swallow ``author``, and the ``token`` +#: alternatives all require a credential-ish prefix or suffix -- a bare +#: ``token`` is a secret, but ``input_tokens`` is a number we want to keep. +_SECRET_KEY = re.compile( + r"(api[_-]?key|secret|password|passwd|credential|bearer|" + r"authorization|auth[_-]|\bauth\b|" + r"(?:access|refresh|id|api|auth|session|csrf|jwt)[_-]?token|\btokens?\b|" + r"private[_-]?key|access[_-]?key|session[_-]?key|client[_-]?secret)", + re.IGNORECASE, +) + +#: Checked before :data:`_SECRET_KEY` and wins: these are token *counts*, which +#: are among the most useful things captured provenance holds. Redacting them +#: would be a silent data-quality bug rather than a safety win. +_TOKEN_COUNT_KEY = re.compile( + r"^(?:(?:input|output|total|prompt|completion|reasoning|cache\w*|max|num|n)[_-]?tokens?" + r"|tokens?[_-]?(?:count|used|total)" + r"|tokens)$", + re.IGNORECASE, +) + +#: Literal shapes that look like credentials wherever they appear in text. +_SECRET_VALUE_PATTERNS = [ + re.compile(r"sk-[A-Za-z0-9_\-]{16,}"), # OpenAI-style + re.compile(r"sk-ant-[A-Za-z0-9_\-]{16,}"), # Anthropic + re.compile(r"gh[pousr]_[A-Za-z0-9]{16,}"), # GitHub + re.compile(r"AKIA[0-9A-Z]{16}"), # AWS access key id + re.compile(r"AIza[0-9A-Za-z_\-]{20,}"), # Google + re.compile(r"xox[baprs]-[A-Za-z0-9\-]{10,}"), # Slack + re.compile(r"eyJ[A-Za-z0-9_\-]{10,}\.[A-Za-z0-9_\-]{10,}\.[A-Za-z0-9_\-]{10,}"), # JWT + re.compile(r"-----BEGIN [A-Z ]*PRIVATE KEY-----"), +] + +#: Tool-input keys that carry file bodies rather than parameters. +_CONTENT_KEYS = {"content", "new_string", "old_string", "new_str", "old_str", "file_text", "text", "patch", "diff"} + +_MAX_DEPTH = 6 +_MAX_ITEMS = 100 + + +def redact_text(text: str) -> str: + """Mask credential-shaped substrings inside free text.""" + for pattern in _SECRET_VALUE_PATTERNS: + text = pattern.sub(REDACTED, text) + return text + + +def _summarize(text: str) -> dict[str, Any]: + """Replace a body with its shape: enough to trace, too little to leak.""" + return { + "_summary": True, + "chars": len(text), + "lines": text.count("\n") + 1 if text else 0, + "sha256_16": content_digest(text), + "preview": text[:200], + } + + +class Sanitizer: + """Applies redaction, truncation, and JSON-coercion under a `Config`.""" + + def __init__(self, config: Config): + self.config = config + + def value(self, obj: Any, *, key: str | None = None, depth: int = 0) -> Any: + """Sanitize an arbitrary value for storage.""" + cfg = self.config + + if ( + cfg.redact + and key + and not _TOKEN_COUNT_KEY.match(key) + and _SECRET_KEY.search(key) + and isinstance(obj, (str, int, float)) + ): + return REDACTED + + if obj is None or isinstance(obj, (bool, int, float)): + return obj + + if isinstance(obj, str): + return self._string(obj, key=key) + + if depth >= _MAX_DEPTH: + return f"" + + if isinstance(obj, dict): + out: dict[str, Any] = {} + for i, (k, v) in enumerate(obj.items()): + if i >= _MAX_ITEMS: + out["_truncated_keys"] = len(obj) - _MAX_ITEMS + break + out[str(k)] = self.value(v, key=str(k), depth=depth + 1) + return out + + if isinstance(obj, (list, tuple, set)): + items = list(obj) + out_list = [self.value(v, key=key, depth=depth + 1) for v in items[:_MAX_ITEMS]] + if len(items) > _MAX_ITEMS: + out_list.append(f"<{len(items) - _MAX_ITEMS} more items>") + return out_list + + # Anything else: best-effort JSON, else repr. + try: + json.dumps(obj) + return obj + except (TypeError, ValueError): + return self._string(repr(obj), key=key) + + def _string(self, text: str, *, key: str | None) -> Any: + cfg = self.config + + is_content = key is not None and key.lower() in _CONTENT_KEYS + if is_content: + if cfg.content_mode == CONTENT_NONE: + return {"_summary": True, "chars": len(text), "omitted": True} + if cfg.content_mode == CONTENT_SUMMARY: + summary = _summarize(text) + if cfg.redact: + summary["preview"] = redact_text(summary["preview"]) + return summary + # CONTENT_FULL falls through to the normal truncation path. + assert cfg.content_mode == CONTENT_FULL + + if cfg.redact: + text = redact_text(text) + + if len(text) > cfg.max_str: + kept = cfg.max_str + return text[:kept] + f"…" + return text + + def mapping(self, obj: Any) -> dict[str, Any]: + """Sanitize a value that must end up as a dict (``used``/``generated``).""" + result = self.value(obj) + if isinstance(result, dict): + return result + return {"value": result} diff --git a/src/flowcept/agents/harness/state.py b/src/flowcept/agents/harness/state.py new file mode 100644 index 00000000..d7cd0bf1 --- /dev/null +++ b/src/flowcept/agents/harness/state.py @@ -0,0 +1,162 @@ +"""Per-session state shared across independent hook processes. + +A harness fires each lifecycle event in its own process, and parallel tool +calls mean several of those processes can run at the same instant. State is +therefore a small JSON file guarded by an advisory lock, always read-modify- +written inside the lock so concurrent ``PostToolUse`` hooks cannot clobber each +other. +""" + +from __future__ import annotations + +import json +import os +import time +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path +from typing import Any + +try: # POSIX + import fcntl + + _HAVE_FCNTL = True +except ImportError: # pragma: no cover - Windows + fcntl = None # type: ignore[assignment] + _HAVE_FCNTL = False + +try: # Windows + import msvcrt + + _HAVE_MSVCRT = True +except ImportError: + msvcrt = None # type: ignore[assignment] + _HAVE_MSVCRT = False + + +#: Pending tool calls older than this are dropped: their PostToolUse never +#: arrived (the harness crashed, or the call was interrupted). +PENDING_TTL_SECONDS = 6 * 60 * 60 +MAX_PENDING = 512 + + +@contextmanager +def file_lock(path: Path, timeout: float = 5.0) -> Iterator[None]: + """Acquire an exclusive advisory lock on ``path`` (``path`` need not exist).""" + path.parent.mkdir(parents=True, exist_ok=True) + lock_path = path.with_suffix(path.suffix + ".lock") + fd = os.open(str(lock_path), os.O_RDWR | os.O_CREAT, 0o600) + deadline = time.monotonic() + timeout + try: + while True: + try: + if _HAVE_FCNTL: + fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + elif _HAVE_MSVCRT: # pragma: no cover - Windows + msvcrt.locking(fd, msvcrt.LK_NBLCK, 1) + break + except OSError: + if time.monotonic() >= deadline: + # Never block the harness on a stuck lock: proceed unlocked + # and accept the small risk of a lost concurrent update. + break + time.sleep(0.01) + yield + finally: + try: + if _HAVE_FCNTL: + fcntl.flock(fd, fcntl.LOCK_UN) + elif _HAVE_MSVCRT: # pragma: no cover - Windows + msvcrt.locking(fd, msvcrt.LK_UNLCK, 1) + except OSError: + pass + os.close(fd) + + +class SessionState: + """Mutable view of one session's state file.""" + + def __init__(self, path: Path, data: dict[str, Any]): + self.path = path + self.data = data + + # -- lifecycle --------------------------------------------------------- + + @staticmethod + def _read(path: Path) -> dict[str, Any]: + try: + with path.open("r", encoding="utf-8") as fh: + data = json.load(fh) + return data if isinstance(data, dict) else {} + except (FileNotFoundError, json.JSONDecodeError, OSError): + return {} + + def save(self) -> None: + """Atomically persist state (write to a temp file, then rename).""" + self.path.parent.mkdir(parents=True, exist_ok=True) + tmp = self.path.with_suffix(f".{os.getpid()}.tmp") + with tmp.open("w", encoding="utf-8") as fh: + json.dump(self.data, fh) + os.replace(tmp, self.path) + + # -- accessors --------------------------------------------------------- + + def get(self, key: str, default: Any = None) -> Any: + """Return the value stored under ``key``, or ``default`` if absent.""" + return self.data.get(key, default) + + def set(self, key: str, value: Any) -> None: + """Store ``value`` under ``key``.""" + self.data[key] = value + + def setdefault(self, key: str, value: Any) -> Any: + """Store ``value`` under ``key`` only if absent, returning the stored value.""" + return self.data.setdefault(key, value) + + def mark_once(self, key: str) -> bool: + """Return True the first time ``key`` is claimed, False afterwards. + + Used to emit the workflow and agent records exactly once per session + even though several processes race to be "first". + """ + flags = self.data.setdefault("_once", {}) + if flags.get(key): + return False + flags[key] = True + return True + + # -- pending tool calls ------------------------------------------------ + + def add_pending(self, key: str, payload: dict[str, Any]) -> None: + """Record an in-flight tool call under ``key``, pruning stale entries.""" + pending = self.data.setdefault("pending_tools", {}) + payload.setdefault("recorded_at", time.time()) + pending[key] = payload + self._prune_pending(pending) + + def pop_pending(self, key: str) -> dict[str, Any] | None: + """Remove and return the pending tool call under ``key``, if any.""" + pending = self.data.setdefault("pending_tools", {}) + value = pending.pop(key, None) + self._prune_pending(pending) + return value + + @staticmethod + def _prune_pending(pending: dict[str, Any]) -> None: + now = time.time() + stale = [k for k, v in pending.items() if now - float(v.get("recorded_at", now)) > PENDING_TTL_SECONDS] + for key in stale: + pending.pop(key, None) + if len(pending) > MAX_PENDING: + ordered = sorted(pending.items(), key=lambda kv: kv[1].get("recorded_at", 0)) + for key, _ in ordered[: len(pending) - MAX_PENDING]: + pending.pop(key, None) + + +@contextmanager +def session_state(path: Path) -> Iterator[SessionState]: + """Open session state for read-modify-write under an exclusive lock.""" + with file_lock(path): + state = SessionState(path, SessionState._read(path)) + yield state + state.save() diff --git a/src/flowcept/agents/harness/tracer.py b/src/flowcept/agents/harness/tracer.py new file mode 100644 index 00000000..adf5a2ff --- /dev/null +++ b/src/flowcept/agents/harness/tracer.py @@ -0,0 +1,318 @@ +"""A session you hold onto, for agents that run inside one process. + +Hook adapters get a fresh process per event and rebuild their context from the +session file each time. An SDK does not: it has a run object, a callback +protocol, and a lifetime. :class:`SessionTracer` is the in-process counterpart — +it holds the sticky facts about the run (harness, model, cwd) so callbacks only +have to say what happened, not who it happened to. + + tracer = SessionTracer("my_agent", model="claude-opus-5") + with tracer: + tracer.prompt("summarize the repo") + with tracer.tool("read_file", {"path": "README.md"}) as call: + call.result({"bytes": 4096}) + tracer.turn_end(response="...", usage={"input_tokens": 900}) + +State still round-trips through the session file on every event, exactly as it +does for hooks. That costs a locked read-modify-write per callback, and buys +the thing that matters more: a run killed mid-flight leaves provenance that is +readable and correctly attributed rather than nothing at all. +""" + +from __future__ import annotations + +import os +import threading +import time +import uuid +from typing import Any + +from .config import Config +from .events import HarnessEvent +from .recorder import Recorder +from .vocab import EventKind + + +class SessionTracer: + """Records one SDK-driven agent run as a Flowcept session. + + Parameters + ---------- + harness: + Identifies the SDK in the provenance, e.g. ``claude_agent_sdk``. + session_id: + The SDK's own conversation/thread id when it has one. All provenance + ids derive from it, so passing the SDK's id is what lets a resumed run + land in the same workflow. A random one is generated otherwise. + model, cwd, project_dir: + Sticky context replayed onto every event, so callbacks that only know + about a tool call still produce fully attributed records. + source_agent_id: + A task or agent id from another capture system (e.g. a + framework-emitted task that launched this run). Stamped as + ``source_agent_id`` on every turn/tool/LLM task the run emits. Falls + back to ``FLOWCEPT_HARNESS_SOURCE_AGENT_ID`` via the config. + """ + + def __init__( + self, + harness: str, + session_id: str | None = None, + *, + config: Config | None = None, + model: str | None = None, + cwd: str | None = None, + project_dir: str | None = None, + source_agent_id: str | None = None, + recorder: Recorder | None = None, + ): + self.harness = harness + self.session_id = session_id or uuid.uuid4().hex + self.recorder = recorder or Recorder(config) + self.model = model + self.cwd = cwd or _safe_cwd() + self.project_dir = project_dir or os.environ.get("FLOWCEPT_HARNESS_PROJECT_DIR") or self.cwd + self.source_agent_id = source_agent_id + self._started = False + self._ended = False + # SDK callbacks can fire from a worker thread and from the main thread + # in the same run; the session file's lock excludes other processes, + # this excludes ourselves. + self._lock = threading.Lock() + + # -- low level ----------------------------------------------------------- + + def event(self, kind: str, **fields: Any) -> list[dict[str, Any]]: + """Record one event, filling in the run's sticky context.""" + fields.setdefault("model", self.model) + fields.setdefault("cwd", self.cwd) + fields.setdefault("project_dir", self.project_dir) + fields.setdefault("source_agent_id", self.source_agent_id) + event = HarnessEvent(kind=kind, harness=self.harness, session_id=self.session_id, **fields) + with self._lock: + return self.recorder.record(event) + + # -- session lifecycle --------------------------------------------------- + + def start(self, *, source: str = "sdk", model: str | None = None, **fields: Any): + """Open the session. Idempotent, and optional — any event opens it.""" + if model: + self.model = model + if self._started: + return [] + self._started = True + return self.event(EventKind.SESSION_START, source=source, **fields) + + def end(self, *, source: str = "completed", **fields: Any): + """Close the session and write its final totals. Idempotent.""" + if self._ended: + return [] + self._ended = True + return self.event(EventKind.SESSION_END, source=source, **fields) + + def __enter__(self) -> SessionTracer: + """Start the session and return the tracer.""" + self.start() + return self + + def __exit__(self, exc_type, exc, tb) -> bool: + """End the session, marking it as errored if an exception escaped.""" + self.end(source="error" if exc_type else "completed") + return False + + # -- turns --------------------------------------------------------------- + + def prompt(self, text: str | None = None, *, prompt_id: str | None = None, **fields: Any): + """Begin a turn. Any tool calls until :meth:`turn_end` hang off it.""" + return self.event(EventKind.PROMPT, prompt=text, prompt_id=prompt_id, **fields) + + def turn_end( + self, + response: str | None = None, + *, + usage: dict[str, Any] | None = None, + **fields: Any, + ): + """Close the open turn with the assistant's answer.""" + return self.event(EventKind.TURN_END, response=response, usage=usage, **fields) + + def llm_call( + self, + *, + model: str | None = None, + prompt: str | None = None, + response: str | None = None, + usage: dict[str, Any] | None = None, + call_id: str | None = None, + started_at: float | None = None, + error: str | None = None, + **fields: Any, + ): + """Record a single model request, nested under the open turn. + + This is the granularity a hook cannot see. An SDK can, so SDK-captured + sessions carry both: turn-level invocations and the calls inside them. + """ + return self.event( + EventKind.LLM_CALL, + model=model or self.model, + prompt=prompt, + response=response, + usage=usage, + call_id=call_id, + started_at=started_at, + error=error, + **fields, + ) + + # -- tools --------------------------------------------------------------- + + def tool_start( + self, + name: str, + tool_input: Any = None, + *, + tool_use_id: str | None = None, + agent_ref: str | None = None, + **fields: Any, + ) -> str: + """Note that a tool is about to run; returns the id to close it with.""" + tool_use_id = tool_use_id or uuid.uuid4().hex + self.event( + EventKind.TOOL_PRE, + tool_name=name, + tool_input=tool_input, + tool_use_id=tool_use_id, + agent_ref=agent_ref, + **fields, + ) + return tool_use_id + + def tool_end( + self, + tool_use_id: str, + *, + name: str | None = None, + tool_response: Any = None, + error: str | None = None, + agent_ref: str | None = None, + started_at: float | None = None, + **fields: Any, + ): + """Close a tool call opened by :meth:`tool_start`. + + Unpaired calls are fine: the recorder falls back to what this event + carries, so a tool the SDK only reports on completion still lands. + """ + return self.event( + EventKind.TOOL_ERROR if error else EventKind.TOOL_POST, + tool_name=name, + tool_use_id=tool_use_id, + tool_response=tool_response, + error=error, + agent_ref=agent_ref, + started_at=started_at, + **fields, + ) + + def tool(self, name: str, tool_input: Any = None, **fields: Any) -> _ToolCall: + """Context manager form: records the call and any exception it raises. + + with tracer.tool("run_tests", {"suite": "unit"}) as call: + call.result(run_tests()) + """ + return _ToolCall(self, name, tool_input, fields) + + # -- subagents ----------------------------------------------------------- + + def subagent_start( + self, + agent_name: str, + *, + agent_ref: str | None = None, + prompt: str | None = None, + **fields: Any, + ) -> str: + """Open a nested workflow for a subagent; returns its ref.""" + agent_ref = agent_ref or uuid.uuid4().hex + self.event( + EventKind.SUBAGENT_START, + agent_name=agent_name, + agent_ref=agent_ref, + prompt=prompt, + **fields, + ) + return agent_ref + + def subagent_stop( + self, + agent_ref: str, + *, + agent_name: str | None = None, + response: str | None = None, + error: str | None = None, + **fields: Any, + ): + """Close a subagent's nested workflow.""" + return self.event( + EventKind.SUBAGENT_STOP, + agent_ref=agent_ref, + agent_name=agent_name, + response=response, + error=error, + **fields, + ) + + # -- lifecycle ----------------------------------------------------------- + + def notify(self, message: str, **fields: Any): + """Record a notification event carrying ``message``.""" + return self.event(EventKind.NOTIFICATION, message=message, **fields) + + def compact(self, *, source: str | None = None, **fields: Any): + """Record a context-compaction event.""" + return self.event(EventKind.COMPACT, source=source, **fields) + + +class _ToolCall: + """The context manager returned by :meth:`SessionTracer.tool`.""" + + __slots__ = ("_fields", "_id", "_input", "_name", "_response", "_started", "_tracer") + + def __init__(self, tracer: SessionTracer, name: str, tool_input: Any, fields: dict[str, Any]): + self._tracer = tracer + self._name = name + self._input = tool_input + self._fields = fields + self._id: str | None = None + self._response: Any = None + self._started = 0.0 + + def __enter__(self) -> _ToolCall: + self._started = time.time() + self._id = self._tracer.tool_start(self._name, self._input, **self._fields) + return self + + def result(self, value: Any) -> Any: + """Record what the tool returned. Returns it, so it can wrap a call.""" + self._response = value + return value + + def __exit__(self, exc_type, exc, tb) -> bool: + self._tracer.tool_end( + self._id or "", + name=self._name, + tool_response=self._response, + error=f"{exc_type.__name__}: {exc}" if exc_type else None, + started_at=self._started, + agent_ref=self._fields.get("agent_ref"), + ) + return False + + +def _safe_cwd() -> str | None: + try: + return os.getcwd() + except OSError: + # A deleted working directory is not a reason to lose the whole run. + return None diff --git a/src/flowcept/agents/harness/vocab.py b/src/flowcept/agents/harness/vocab.py new file mode 100644 index 00000000..757cd138 --- /dev/null +++ b/src/flowcept/agents/harness/vocab.py @@ -0,0 +1,83 @@ +"""Vocabulary used when mapping harness activity onto PROV-AGENT. + +PROV-AGENT (arXiv:2508.02866) is the W3C PROV extension Flowcept uses for +agentic workflows. It defines two activity classes that map cleanly onto what a +coding harness does, and Flowcept's UI and query layer already understand them: + +``ai_model_invocation`` + One prompt -> one response. We record this at *turn* granularity for + hook-based harnesses (a hook cannot see individual API calls) and at + *call* granularity for SDK and OpenTelemetry adapters, which can. + ``custom_metadata.granularity`` says which. + +``agent_tool`` + One tool execution by the agent — a ``Bash`` run, an ``Edit``, an MCP tool. + +Everything else a harness emits (compaction, notifications, permission +decisions) is lifecycle context rather than a PROV activity class, so it is +recorded with the ``harness_event`` subtype and kept out of the two classes +above so dataflow queries stay clean. +""" + +from __future__ import annotations + +# --- PROV-AGENT activity subtypes (must match flowcept.commons.vocabulary) --- +AI_MODEL_INVOCATION = "ai_model_invocation" +AGENT_TOOL = "agent_tool" + +# --- flowcept-harness extensions -------------------------------------------- +HARNESS_EVENT = "harness_event" +"""Lifecycle event with no dataflow of its own (compaction, notification).""" + +AGENT_SESSION = "agent_session" +"""Workflow subtype for one interactive harness session.""" + +SUBAGENT_SESSION = "subagent_session" +"""Workflow subtype for a subagent nested under a session.""" + +# --- Statuses (must match flowcept.commons.vocabulary.Status) --------------- +STATUS_RUNNING = "RUNNING" +STATUS_FINISHED = "FINISHED" +STATUS_ERROR = "ERROR" +STATUS_UNKNOWN = "UNKNOWN" + +# --- Record types (the "type" discriminator Flowcept's inserter reads) ------ +TYPE_TASK = "task" +TYPE_WORKFLOW = "workflow" +TYPE_AGENT = "agent" + +ADAPTER_ID = "flowcept.agents.harness" + + +#: Normalized event kinds the recorder understands. Adapters translate their +#: harness's native event names into these. +class EventKind: + """Harness-independent lifecycle events.""" + + SESSION_START = "session_start" + SESSION_END = "session_end" + PROMPT = "prompt" + TURN_END = "turn_end" + TOOL_PRE = "tool_pre" + TOOL_POST = "tool_post" + TOOL_ERROR = "tool_error" + LLM_CALL = "llm_call" + SUBAGENT_START = "subagent_start" + SUBAGENT_STOP = "subagent_stop" + NOTIFICATION = "notification" + COMPACT = "compact" + + ALL = ( + SESSION_START, + SESSION_END, + PROMPT, + TURN_END, + TOOL_PRE, + TOOL_POST, + TOOL_ERROR, + LLM_CALL, + SUBAGENT_START, + SUBAGENT_STOP, + NOTIFICATION, + COMPACT, + ) diff --git a/src/flowcept/agents/langchain/langchain_plugin.py b/src/flowcept/agents/langchain/langchain_plugin.py new file mode 100644 index 00000000..df69908e --- /dev/null +++ b/src/flowcept/agents/langchain/langchain_plugin.py @@ -0,0 +1,401 @@ +"""LangChain and LangGraph wrapper. + +LangChain's callback protocol is already the right shape for provenance: every +run reports a start, an end, its own ``run_id``, and its ``parent_run_id``. +LangGraph uses the same protocol, so one handler covers both. + + from flowcept.agents.langchain.langchain_plugin import FlowceptCallbackHandler + + handler = FlowceptCallbackHandler(session_id="thread-42") + graph.invoke({"messages": [...]}, config={"callbacks": [handler]}) + +What each run becomes: + +===================== ================================================== +root chain / graph the turn — its inputs are the prompt, its outputs + the response +nested chain / node structure only, no record of its own +LLM / chat model ``ai_model_invocation`` at call granularity +tool ``agent_tool`` +retriever ``agent_tool`` (a retrieval is a tool execution) +===================== ================================================== + +Nested chains are deliberately not recorded. LangGraph emits a run per node, +per branch, and per internal ``RunnableSequence``, and turning all of that into +tasks would bury the model calls and tool calls that actually describe what the +agent did. The graph's structure is still recoverable: every recorded task +carries its enclosing turn. + +Duck-typed rather than subclassing ``BaseCallbackHandler`` -- nothing here +imports langchain, so it works against any 0.1+ version and imports without +langchain installed. The ``ignore_*`` and ``raise_error`` attributes below are +part of that contract: the callback manager reads them off the handler. +""" + +from __future__ import annotations + +import json +from typing import Any +from uuid import UUID + +from flowcept.agents.harness.config import Config +from flowcept.agents.harness.tracer import SessionTracer + + +class FlowceptCallbackHandler: + """A LangChain callback handler that records Flowcept provenance.""" + + # -- the attributes langchain's callback manager reads off a handler ----- + ignore_llm = False + ignore_chain = False + ignore_agent = False + ignore_retriever = False + ignore_chat_model = False + ignore_retry = True + ignore_custom_event = True + raise_error = False + run_inline = False + + def __init__( + self, + session_id: str | None = None, + *, + config: Config | None = None, + harness: str = "langchain", + model: str | None = None, + tracer: SessionTracer | None = None, + ): + self.tracer = tracer or SessionTracer(harness, session_id, config=config, model=model) + self.tracer.start() + #: The run that owns the current turn; ``None`` between turns. + self._turn_run: str | None = None + #: run_id -> tool name, so an end callback can name what it closes. + self._tools: dict[str, str] = {} + #: run_id -> model name, likewise for LLM runs. + self._models: dict[str, str | None] = {} + self._prompts: dict[str, str | None] = {} + + # -- chains / graphs ----------------------------------------------------- + + def on_chain_start( + self, + serialized: dict[str, Any] | None, + inputs: Any, + *, + run_id: UUID | None = None, + parent_run_id: UUID | None = None, + **kwargs: Any, + ) -> None: + """Open a turn when a root chain or graph run starts.""" + if parent_run_id is not None or self._turn_run is not None: + return # nested run: structure only + self._open_turn(_key(run_id), _as_text(_unwrap_input(inputs))) + + def on_chain_end(self, outputs: Any, *, run_id: UUID | None = None, **kwargs: Any) -> None: + """Close the turn with the root chain's outputs as the response.""" + if self._turn_run != _key(run_id): + return + self._close_turn(response=_as_text(_unwrap_output(outputs))) + + def on_chain_error(self, error: BaseException, *, run_id: UUID | None = None, **kwargs: Any) -> None: + """Close the turn with the root chain's error.""" + if self._turn_run != _key(run_id): + return + self._close_turn(error=_error_text(error)) + + # -- models -------------------------------------------------------------- + + def on_llm_start( + self, + serialized: dict[str, Any] | None, + prompts: list[str] | None, + *, + run_id: UUID | None = None, + parent_run_id: UUID | None = None, + invocation_params: dict[str, Any] | None = None, + **kwargs: Any, + ) -> None: + """Remember the model and prompt of a starting LLM run.""" + key = _key(run_id) + model = _model_name(serialized, invocation_params, kwargs) + self._models[key] = model + self._prompts[key] = "\n\n".join(p for p in (prompts or []) if isinstance(p, str)) or None + if parent_run_id is None and self._turn_run is None: + # A model invoked directly, with no chain around it: that call is + # the whole turn. + self._open_turn(key, self._prompts[key]) + + def on_chat_model_start( + self, + serialized: dict[str, Any] | None, + messages: Any, + *, + run_id: UUID | None = None, + parent_run_id: UUID | None = None, + invocation_params: dict[str, Any] | None = None, + **kwargs: Any, + ) -> None: + """Handle a starting chat-model run as an LLM run with flattened messages.""" + self.on_llm_start( + serialized, + [_messages_text(messages)] if messages is not None else None, + run_id=run_id, + parent_run_id=parent_run_id, + invocation_params=invocation_params, + **kwargs, + ) + + def on_llm_end(self, response: Any, *, run_id: UUID | None = None, **kwargs: Any) -> None: + """Record the completed LLM run as an ``ai_model_invocation`` task.""" + key = _key(run_id) + model, usage = _llm_result_details(response) + self.tracer.llm_call( + model=model or self._models.pop(key, None), + prompt=self._prompts.pop(key, None), + response=_generations_text(response), + usage=usage, + call_id=key, + ) + self._models.pop(key, None) + if self._turn_run == key: + self._close_turn(response=_generations_text(response)) + + def on_llm_error(self, error: BaseException, *, run_id: UUID | None = None, **kwargs: Any) -> None: + """Record the failed LLM run with its error text.""" + key = _key(run_id) + self.tracer.llm_call( + model=self._models.pop(key, None), + prompt=self._prompts.pop(key, None), + call_id=key, + error=_error_text(error), + ) + if self._turn_run == key: + self._close_turn(error=_error_text(error)) + + # -- tools and retrievers ------------------------------------------------ + + def on_tool_start( + self, + serialized: dict[str, Any] | None, + input_str: str | None, + *, + run_id: UUID | None = None, + inputs: dict[str, Any] | None = None, + **kwargs: Any, + ) -> None: + """Record the start of a tool run as an ``agent_tool`` task.""" + key = _key(run_id) + name = (serialized or {}).get("name") or "tool" + self._tools[key] = name + self.tracer.tool_start(name, inputs if inputs is not None else input_str, tool_use_id=key) + + def on_tool_end(self, output: Any, *, run_id: UUID | None = None, **kwargs: Any) -> None: + """Close the tool run's task with its output.""" + key = _key(run_id) + self.tracer.tool_end(key, name=self._tools.pop(key, None), tool_response=_jsonable(output)) + + def on_tool_error(self, error: BaseException, *, run_id: UUID | None = None, **kwargs: Any) -> None: + """Close the tool run's task with its error.""" + key = _key(run_id) + self.tracer.tool_end(key, name=self._tools.pop(key, None), error=_error_text(error)) + + def on_retriever_start( + self, + serialized: dict[str, Any] | None, + query: str | None, + *, + run_id: UUID | None = None, + **kwargs: Any, + ) -> None: + """Record the start of a retriever run as an ``agent_tool`` task.""" + key = _key(run_id) + name = (serialized or {}).get("name") or "retriever" + self._tools[key] = name + self.tracer.tool_start(name, {"query": query}, tool_use_id=key) + + def on_retriever_end(self, documents: Any, *, run_id: UUID | None = None, **kwargs: Any) -> None: + """Close the retriever run's task with the retrieved documents.""" + key = _key(run_id) + self.tracer.tool_end( + key, + name=self._tools.pop(key, None), + tool_response={"documents": [_document(d) for d in documents or []]}, + ) + + def on_retriever_error(self, error: BaseException, *, run_id: UUID | None = None, **kwargs: Any) -> None: + """Close the retriever run's task with its error.""" + key = _key(run_id) + self.tracer.tool_end(key, name=self._tools.pop(key, None), error=_error_text(error)) + + # -- agents -------------------------------------------------------------- + + def on_agent_action(self, action: Any, **kwargs: Any) -> None: + """No record: the tool callbacks already cover the action itself.""" + + def on_agent_finish(self, finish: Any, **kwargs: Any) -> None: + """No record: the enclosing chain's end closes the turn.""" + + def on_text(self, text: str, **kwargs: Any) -> None: + """No record: intermediate text is not an activity.""" + + # -- teardown ------------------------------------------------------------ + + def close(self, *, error: str | None = None) -> None: + """Close open runs and the session. Safe to call more than once.""" + for key, name in list(self._tools.items()): + self.tracer.tool_end(key, name=name, error="never returned a result") + self._tools.clear() + if self._turn_run is not None: + self._close_turn(error=error) + self.tracer.end(source="error" if error else "completed") + + def __enter__(self) -> FlowceptCallbackHandler: + """Return the handler itself for use as a context manager.""" + return self + + def __exit__(self, exc_type, exc, tb) -> bool: + """Close the session, recording the in-flight exception if any.""" + self.close(error=f"{exc_type.__name__}: {exc}" if exc_type else None) + return False + + # -- turns --------------------------------------------------------------- + + def _open_turn(self, key: str, prompt: str | None) -> None: + self._turn_run = key + self.tracer.prompt(prompt, prompt_id=key) + + def _close_turn(self, *, response: str | None = None, error: str | None = None) -> None: + self._turn_run = None + self.tracer.turn_end(response=response, error=error) + + +# -- readers for langchain's loosely-typed callback payloads ------------------ + + +def _key(run_id: Any) -> str: + return str(run_id) if run_id is not None else "root" + + +def _unwrap_input(inputs: Any) -> Any: + """Pull the interesting part out of a chain's input mapping.""" + if isinstance(inputs, dict): + for field in ("input", "question", "messages", "query"): + if field in inputs: + return inputs[field] + return inputs + + +def _unwrap_output(outputs: Any) -> Any: + if isinstance(outputs, dict): + for field in ("output", "answer", "messages", "result"): + if field in outputs: + return outputs[field] + return outputs + + +def _model_name(serialized: Any, invocation_params: Any, kwargs: dict[str, Any]) -> str | None: + for source in (invocation_params, kwargs.get("metadata"), serialized): + if isinstance(source, dict): + for field in ("model", "model_name", "model_id", "ls_model_name"): + value = source.get(field) + if isinstance(value, str): + return value + if isinstance(serialized, dict): + # Fall back to the class the callback came from, e.g. ChatAnthropic. + identifier = serialized.get("id") + if isinstance(identifier, list) and identifier: + return str(identifier[-1]) + return None + + +def _llm_result_details(response: Any) -> tuple[str | None, dict[str, Any] | None]: + output = getattr(response, "llm_output", None) + model = None + usage = None + if isinstance(output, dict): + model = output.get("model_name") or output.get("model") + for field in ("token_usage", "usage", "usage_metadata"): + candidate = output.get(field) + if isinstance(candidate, dict): + usage = candidate + break + if usage is None: + # Newer versions carry usage on the generation's message instead. + message = _first_generation_attribute(response, "message") + candidate = getattr(message, "usage_metadata", None) + if isinstance(candidate, dict): + usage = candidate + return (model if isinstance(model, str) else None), usage + + +def _generations_text(response: Any) -> str | None: + text = _first_generation_attribute(response, "text") + if isinstance(text, str) and text: + return text + message = _first_generation_attribute(response, "message") + content = getattr(message, "content", None) + return _as_text(content) if content is not None else None + + +def _first_generation_attribute(response: Any, attribute: str) -> Any: + generations = getattr(response, "generations", None) + if not isinstance(generations, list): + return None + for group in generations: + items = group if isinstance(group, list) else [group] + for item in items: + value = getattr(item, attribute, None) + if value is not None: + return value + return None + + +def _messages_text(messages: Any) -> str: + """Flatten the nested list of chat messages a chat model is started with.""" + parts: list[str] = [] + groups = messages if isinstance(messages, list) else [messages] + for group in groups: + items = group if isinstance(group, list) else [group] + for message in items: + content = getattr(message, "content", message) + role = getattr(message, "type", None) or getattr(message, "role", None) + text = _as_text(content) or "" + parts.append(f"{role}: {text}" if role else text) + return "\n".join(parts) + + +def _document(document: Any) -> dict[str, Any]: + return { + "page_content": getattr(document, "page_content", None), + "metadata": _jsonable(getattr(document, "metadata", None)), + } + + +def _error_text(error: BaseException | Any) -> str: + if isinstance(error, BaseException): + return f"{type(error).__name__}: {error}" + return str(error) + + +def _jsonable(value: Any) -> Any: + if value is None or isinstance(value, (dict, list, str, int, float, bool)): + return value + for attribute in ("model_dump", "dict", "to_dict"): + method = getattr(value, attribute, None) + if callable(method): + try: + result = method() + except Exception: + continue + if isinstance(result, dict): + return result + return repr(value) + + +def _as_text(value: Any) -> str | None: + if value is None or isinstance(value, str): + return value + try: + return json.dumps(_jsonable(value), default=repr) + except (TypeError, ValueError): + return repr(value) diff --git a/src/flowcept/agents/langgraph/langgraph_plugin.py b/src/flowcept/agents/langgraph/langgraph_plugin.py new file mode 100644 index 00000000..6a7b0a00 --- /dev/null +++ b/src/flowcept/agents/langgraph/langgraph_plugin.py @@ -0,0 +1,1508 @@ +# academy_coscientist/plugins/flowcept_langgraph_plugin.py +""" +FlowCept provenance plugin for LangGraph workflows. + +Mirrors the design of FlowceptAcademyPlugin but targets LangGraph/LangChain +instead of Academy. Provenance is captured through LangChain's standard +callback interface, which LangGraph respects for all graph, node, LLM, and +tool executions. + +Provenance hierarchy produced: + WorkflowObject (one per graph.invoke / graph.ainvoke call) + └─ TaskObject subtype=langgraph_node activity_id= + └─ TaskObject subtype=llm_call activity_id= (parent_task_id) + └─ TaskObject subtype=tool_call activity_id= (parent_task_id) + +Key FlowCept fields: + task_id — uuid per event (generated here, not LangChain's run_id) + workflow_id — per graph-run sub-workflow id + campaign_id — from Flowcept.campaign_id + parent_task_id — links LLM/tool tasks to their enclosing node task + group_id — all tasks within one graph run share a group_id + activity_id — node name | model name | tool name + subtype — langgraph_node | llm_call | tool_call | langgraph_graph + used / generated — inputs and outputs + status — FINISHED | ERROR (proper Status enum values) + telemetry_at_start/end — CPU/memory snapshots via TelemetryCapture + +Usage (three lines to wire up): + + from academy_coscientist.plugins.flowcept_langgraph_plugin import FlowceptLangGraphPlugin + + plugin = FlowceptLangGraphPlugin(config={"workflow_name": "my-graph"}) + plugin.start() + result = graph.invoke(state, config={"callbacks": [plugin.callback_handler]}) + plugin.stop() + +Or as a context manager: + + with FlowceptLangGraphPlugin(config={"workflow_name": "my-graph"}) as plugin: + result = graph.invoke(state, config={"callbacks": [plugin.callback_handler]}) +""" + +from __future__ import annotations + +import os +import threading +import time +import uuid +import logging +from contextlib import contextmanager +from typing import Any +from uuid import UUID + +_log = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Provenance overhead timer (identical to flowcept_plugin.py) +# --------------------------------------------------------------------------- + + +class _ProvenanceStats: + """Lightweight thread-safe accumulator for provenance capture timings.""" + + __slots__ = ("_lock", "_counts", "_totals", "_mins", "_maxs", "_raw") + + def __init__(self) -> None: + self._lock: threading.Lock = threading.Lock() + self._counts: dict[str, int] = {} + self._totals: dict[str, float] = {} + self._mins: dict[str, float] = {} + self._maxs: dict[str, float] = {} + self._raw: list[tuple[str, str, float]] = [] + + def record(self, category: str, elapsed: float) -> None: + ts = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) + with self._lock: + if category not in self._counts: + self._counts[category] = 0 + self._totals[category] = 0.0 + self._mins[category] = float("inf") + self._maxs[category] = 0.0 + self._counts[category] += 1 + self._totals[category] += elapsed + if elapsed < self._mins[category]: + self._mins[category] = elapsed + if elapsed > self._maxs[category]: + self._maxs[category] = elapsed + self._raw.append((ts, category, elapsed)) + + def summary(self) -> str: + col = 22 + header = f"{'Category':<{col}} {'N':>7} {'Total(ms)':>11} {'Mean(µs)':>9} {'Min(µs)':>8} {'Max(µs)':>8}" + sep = "-" * len(header) + rows = [header, sep] + with self._lock: + for cat in sorted(self._counts): + n = self._counts[cat] + total = self._totals[cat] + mean = (total / n) if n else 0.0 + mn = self._mins.get(cat, 0.0) + mx = self._maxs.get(cat, 0.0) + rows.append( + f"{cat:<{col}} {n:>7} {total * 1e3:>11.3f} {mean * 1e6:>9.1f} {mn * 1e6:>8.1f} {mx * 1e6:>8.1f}" + ) + return "\n".join(rows) + + def to_csv(self, path: str, workflow_id: str | None = None) -> None: + import csv + + write_header = not os.path.exists(path) + with self._lock: + raw_snapshot = list(self._raw) + wf = workflow_id or "" + rows = [ + { + "timestamp_utc": ts, + "workflow_id": wf, + "category": cat, + "elapsed_us": round(elapsed * 1e6, 3), + } + for ts, cat, elapsed in raw_snapshot + ] + fieldnames = ["timestamp_utc", "workflow_id", "category", "elapsed_us"] + with open(path, "a", newline="", encoding="utf-8") as fh: + writer = csv.DictWriter(fh, fieldnames=fieldnames) + if write_header: + writer.writeheader() + writer.writerows(rows) + + +# --------------------------------------------------------------------------- +# Interceptor wrapper (same pattern as AcademyInterceptor) +# --------------------------------------------------------------------------- + + +class LangGraphInterceptor: + """Manages a BaseInterceptor directly for LangGraph provenance (Dask-style).""" + + def __init__(self) -> None: + self._interceptor = None + self._workflow_id: str | None = None + self._campaign_id: str | None = None + + def start(self, workflow_name: str, campaign_id: str | None = None) -> None: + """Start the interceptor and emit the top-level WorkflowObject.""" + from flowcept.flowceptor.adapters.base_interceptor import BaseInterceptor + from flowcept.commons.flowcept_dataclasses.workflow_object import WorkflowObject + + self._workflow_id = str(uuid.uuid4()) + self._campaign_id = campaign_id or str(uuid.uuid4()) + + self._interceptor = BaseInterceptor(kind="langgraph") + self._interceptor.start( + bundle_exec_id=self._workflow_id, + check_safe_stops=False, + ) + + wf = WorkflowObject() + wf.workflow_id = self._workflow_id + wf.campaign_id = self._campaign_id + wf.name = workflow_name + self._interceptor.send_workflow_message(wf) + + def stop(self) -> None: + """Stop the interceptor, flushing any buffered provenance.""" + if self._interceptor is None: + return + try: + self._interceptor.stop(check_safe_stops=False) + except Exception as e: + _log.warning("Interceptor stop error: %r", e) + self._interceptor = None + + @property + def telemetry_capture(self): + """Return the interceptor's TelemetryCapture, or None when not started.""" + return self._interceptor.telemetry_capture if self._interceptor else None + + def send_graph_workflow(self, graph_name: str, group_id: str) -> str: + """Emit a WorkflowObject for one graph invocation; return its workflow_id.""" + if self._interceptor is None: + return str(uuid.uuid4()) + from flowcept.commons.flowcept_dataclasses.workflow_object import WorkflowObject + + wf = WorkflowObject() + wf.workflow_id = str(uuid.uuid4()) + wf.name = graph_name + wf.campaign_id = self._campaign_id + wf.parent_workflow_id = self._workflow_id + wf.custom_metadata = {"group_id": group_id, "graph_name": graph_name} + self._interceptor.send_workflow_message(wf) + return wf.workflow_id + + def intercept_task(self, task_dict: dict) -> None: + """Enrich a task dict with ids and status, then send it to the interceptor.""" + if self._interceptor is None: + return + from flowcept.commons.flowcept_dataclasses.task_object import TaskObject + from flowcept.commons.vocabulary import Status + + task_dict.setdefault("task_id", str(uuid.uuid4())) + task_dict.setdefault("workflow_id", self._workflow_id) + task_dict.setdefault("campaign_id", self._campaign_id) + + raw = task_dict.get("status", "FINISHED") + if isinstance(raw, str): + try: + task_dict["status"] = Status[raw].value + except KeyError: + task_dict["status"] = Status.FINISHED.value + + TaskObject.enrich_task_dict(task_dict) + self._interceptor.intercept(task_dict) + + +# --------------------------------------------------------------------------- +# LangGraph callback handler +# --------------------------------------------------------------------------- + + +class FlowceptLangGraphCallback: + """ + LangChain BaseCallbackHandler that records provenance to FlowCept. + + Pass an instance to ``graph.invoke`` / ``graph.ainvoke`` via the + ``config={"callbacks": [...]}`` argument. + + This class deliberately avoids subclassing ``BaseCallbackHandler`` at + module import time so LangGraph/LangChain are optional dependencies. + The actual class is built lazily in ``_build_handler_class()`` and + cached in ``_HANDLER_CLASS``. + """ + + def __init__(self, interceptor: LangGraphInterceptor, stats: _ProvenanceStats | None) -> None: + self._interceptor = interceptor + self._stats = stats + # Maps LangChain run_id (UUID) → FlowCept task_id (str) + self._run_to_task: dict[UUID, str] = {} + # Maps LangChain run_id → start timestamp + self._run_start: dict[UUID, float] = {} + # Tracks telemetry snapshots at node start + self._run_tel_start: dict[UUID, Any] = {} + # Maps graph-level run_id → group_id (shared by all tasks in one graph invocation) + self._graph_runs: dict[UUID, str] = {} + # Maps any child run_id → enclosing graph-level run_id (for group_id lookup) + self._run_to_graph_run: dict[UUID, UUID] = {} + # Maps graph-level run_id → Academy agent ID that produced the input data. + # Set when the LangGraph state contains "_source_agent_id". + self._graph_source_agent: dict[UUID, str] = {} + # Buffers the start-time task skeleton (used, custom_metadata, parent_task_id …) + # so on_chain_end / on_llm_end can emit ONE complete record with both + # used (inputs) and generated (outputs) — same pattern as the Academy plugin. + self._run_start_task: dict[UUID, dict] = {} + + # --- helpers --- + + def _is_graph_run( + self, + serialized: dict | None, + tags: list[str] | None, + parent_run_id=None, + ) -> bool: + """Return True when the chain event represents a top-level LangGraph graph invocation. + + Primary signal: parent_run_id is None (no enclosing run → this IS the graph). + Secondary signal: serialized id / tags for belt-and-suspenders detection. + """ + if parent_run_id is None: + return True + id_parts = (serialized or {}).get("id", []) + graph_classes = { + "CompiledStateGraph", + "CompiledGraph", + "Pregel", + "StateGraph", + "MessageGraph", + } + return bool(graph_classes.intersection(id_parts) or any(t.startswith("__pregel_") for t in (tags or []))) + + def _node_name(self, serialized: dict | None, name: str | None) -> str: + """Extract a human-readable node name.""" + if name: + return name + id_parts = (serialized or {}).get("id", []) + return id_parts[-1] if id_parts else "unknown_node" + + def _model_name(self, serialized: dict | None) -> str: + id_parts = (serialized or {}).get("id", []) + kwargs = (serialized or {}).get("kwargs", {}) + return kwargs.get("model_name") or kwargs.get("model") or (id_parts[-1] if id_parts else "unknown_model") + + def _tel(self) -> Any: + tc = self._interceptor.telemetry_capture + return tc.capture() if tc else None + + def _record(self, category: str, elapsed: float) -> None: + if self._stats is not None: + self._stats.record(category, elapsed) + + def _task_id_for(self, run_id: UUID) -> str: + if run_id not in self._run_to_task: + self._run_to_task[run_id] = str(uuid.uuid4()) + return self._run_to_task[run_id] + + def _group_id_for(self, run_id: UUID) -> str | None: + """Return the group_id for the graph invocation enclosing this run.""" + if run_id in self._graph_runs: + return self._graph_runs[run_id] + graph_run_id = self._run_to_graph_run.get(run_id) + if graph_run_id is not None: + return self._graph_runs.get(graph_run_id) + return None + + # --- chain events (graph & node) --- + + def on_chain_start( + self, + serialized: dict | None, + inputs: dict | None, + *, + run_id: UUID, + parent_run_id: UUID | None = None, + tags: list[str] | None = None, + name: str | None = None, + **kwargs: Any, + ) -> None: + """Buffer a task skeleton for a starting graph or node run.""" + t0 = time.perf_counter() + self._run_start[run_id] = time.time() + self._run_tel_start[run_id] = self._tel() + task_id = self._task_id_for(run_id) + serialized = serialized or {} + inputs = inputs or {} + + if self._is_graph_run(serialized, tags, parent_run_id): + # Top-level graph invocation → emit a sub-WorkflowObject for + # hierarchy tracking, and record a group_id shared by all tasks + # in this invocation. Tasks use the global workflow_id (not the + # sub-workflow id) so that campaign_id / workflow_id are uniform. + group_id = str(uuid.uuid4()) + graph_name = name or self._node_name(serialized, name) + self._interceptor.send_graph_workflow(graph_name, group_id) + self._graph_runs[run_id] = group_id + source_agent_id = inputs.get("_source_agent_id") + if source_agent_id: + self._graph_source_agent[run_id] = str(source_agent_id) + + custom: dict[str, Any] = {"tags": tags or [], "graph_name": graph_name} + if source_agent_id: + custom["source_agent_id"] = str(source_agent_id) + # Buffer the skeleton — emitted as ONE complete record in on_chain_end + self._run_start_task[run_id] = { + "task_id": task_id, + "subtype": "langgraph_graph", + "activity_id": graph_name, + "group_id": group_id, + "started_at": self._run_start[run_id], + "used": {"inputs": _safe_clip(inputs)}, + "custom_metadata": custom, + } + else: + # Node or sub-chain execution. + node = self._node_name(serialized, name) + # Find the enclosing graph run for group_id and source_agent_id. + graph_run_id = None + if parent_run_id and parent_run_id in self._graph_runs: + graph_run_id = parent_run_id + elif parent_run_id: + graph_run_id = self._run_to_graph_run.get(parent_run_id) + group_id = self._graph_runs.get(graph_run_id) if graph_run_id else None + source_agent_id = self._graph_source_agent.get(graph_run_id) if graph_run_id else None + + # Record child → graph mapping so on_llm_end / on_tool_end can find group_id + if graph_run_id is not None: + self._run_to_graph_run[run_id] = graph_run_id + + parent_task_id = self._run_to_task.get(parent_run_id) if parent_run_id else None + custom = {"tags": tags or [], "node_name": node} + if source_agent_id: + custom["source_agent_id"] = source_agent_id + # Buffer the skeleton — emitted as ONE complete record in on_chain_end + skeleton: dict[str, Any] = { + "task_id": task_id, + "subtype": "langgraph_node", + "activity_id": node, + "group_id": group_id, + "started_at": self._run_start[run_id], + "used": {"inputs": _safe_clip(inputs)}, + "custom_metadata": custom, + } + if parent_task_id: + skeleton["parent_task_id"] = parent_task_id + self._run_start_task[run_id] = skeleton + + self._record("chain_start", time.perf_counter() - t0) + + def on_chain_end( + self, + outputs: dict, + *, + run_id: UUID, + parent_run_id: UUID | None = None, + **kwargs: Any, + ) -> None: + """Emit the completed graph or node task with outputs and telemetry.""" + t0 = time.perf_counter() + self._run_start.pop(run_id, time.time()) + tel_start = self._run_tel_start.pop(run_id, None) + tel_end = self._tel() + self._run_to_task.pop(run_id, None) + + is_graph = run_id in self._graph_runs + if is_graph: + self._graph_runs.pop(run_id) + self._graph_source_agent.pop(run_id, None) + else: + self._run_to_graph_run.pop(run_id, None) + + # Merge buffered start skeleton with completion data → ONE complete record + task: dict[str, Any] = self._run_start_task.pop(run_id, {}) + task.update( + { + "ended_at": time.time(), + "status": "FINISHED", + "generated": {"outputs": _safe_clip(outputs)}, + } + ) + if tel_start is not None: + task["telemetry_at_start"] = _tel_to_dict(tel_start) + if tel_end is not None: + task["telemetry_at_end"] = _tel_to_dict(tel_end) + + self._interceptor.intercept_task(task) + self._record("chain_end", time.perf_counter() - t0) + + def on_chain_error( + self, + error: BaseException, + *, + run_id: UUID, + parent_run_id: UUID | None = None, + **kwargs: Any, + ) -> None: + """Emit the failed graph or node task with ERROR status.""" + t0 = time.perf_counter() + self._run_start.pop(run_id, time.time()) + tel_start = self._run_tel_start.pop(run_id, None) + self._run_to_task.pop(run_id, None) + is_graph = run_id in self._graph_runs + if is_graph: + self._graph_runs.pop(run_id, None) + self._graph_source_agent.pop(run_id, None) + else: + self._run_to_graph_run.pop(run_id, None) + + task: dict[str, Any] = self._run_start_task.pop(run_id, {}) + task.update( + { + "ended_at": time.time(), + "status": "ERROR", + "stderr": str(error), + } + ) + if tel_start is not None: + task["telemetry_at_start"] = _tel_to_dict(tel_start) + self._interceptor.intercept_task(task) + self._record("chain_error", time.perf_counter() - t0) + + # --- LLM events --- + + def on_llm_start( + self, + serialized: dict | None, + prompts: list[str], + *, + run_id: UUID, + parent_run_id: UUID | None = None, + **kwargs: Any, + ) -> None: + """Buffer an llm_call task skeleton for a starting LLM run.""" + t0 = time.perf_counter() + self._run_start[run_id] = time.time() + self._run_tel_start[run_id] = self._tel() + task_id = self._task_id_for(run_id) + # Propagate enclosing graph run mapping for group_id lookup in on_llm_end + if parent_run_id is not None: + graph_run_id = self._run_to_graph_run.get(parent_run_id) or ( + parent_run_id if parent_run_id in self._graph_runs else None + ) + if graph_run_id is not None: + self._run_to_graph_run[run_id] = graph_run_id + group_id = self._group_id_for(run_id) + parent_task_id = self._run_to_task.get(parent_run_id) if parent_run_id else None + model_name = self._model_name(serialized) + skeleton: dict[str, Any] = { + "task_id": task_id, + "subtype": "llm_call", + "activity_id": model_name, + "group_id": group_id, + "started_at": self._run_start[run_id], + "used": {"prompts": _safe_clip(prompts), "model": model_name}, + "custom_metadata": {"model": model_name}, + } + if parent_task_id: + skeleton["parent_task_id"] = parent_task_id + self._run_start_task[run_id] = skeleton + self._record("llm_start", time.perf_counter() - t0) + + def on_llm_end( + self, + response: Any, + *, + run_id: UUID, + parent_run_id: UUID | None = None, + **kwargs: Any, + ) -> None: + """Emit the completed llm_call task with response text and token usage.""" + t0 = time.perf_counter() + self._run_start.pop(run_id, time.time()) + tel_start = self._run_tel_start.pop(run_id, None) + tel_end = self._tel() + self._run_to_task.pop(run_id, None) + self._run_to_graph_run.pop(run_id, None) + + # Extract token usage and text from LangChain LLMResult + usage: dict = {} + text: str = "" + model_name: str = "unknown" + try: + gen = response.generations + if gen and gen[0]: + text = gen[0][0].text if hasattr(gen[0][0], "text") else str(gen[0][0]) + llm_output = response.llm_output or {} + usage = llm_output.get("token_usage") or llm_output.get("usage") or {} + model_name = llm_output.get("model_name") or llm_output.get("model") or "unknown" + except Exception: + pass + + # Merge buffered start skeleton (includes prompts in used) with response data + task: dict[str, Any] = self._run_start_task.pop(run_id, {}) + task.update( + { + "ended_at": time.time(), + "status": "FINISHED", + "generated": { + "text": text, + "prompt_tokens": usage.get("prompt_tokens"), + "completion_tokens": usage.get("completion_tokens"), + "total_tokens": usage.get("total_tokens"), + "model": model_name, + }, + } + ) + # Update activity_id and model in custom_metadata if model is now known + if model_name != "unknown": + task["activity_id"] = model_name + task.setdefault("custom_metadata", {})["model"] = model_name + if tel_start is not None: + task["telemetry_at_start"] = _tel_to_dict(tel_start) + if tel_end is not None: + task["telemetry_at_end"] = _tel_to_dict(tel_end) + + self._interceptor.intercept_task(task) + self._record("llm_end", time.perf_counter() - t0) + + def on_llm_error( + self, + error: BaseException, + *, + run_id: UUID, + parent_run_id: UUID | None = None, + **kwargs: Any, + ) -> None: + """Emit the failed llm_call task with ERROR status.""" + t0 = time.perf_counter() + self._run_start.pop(run_id, None) + tel_start = self._run_tel_start.pop(run_id, None) + self._run_to_task.pop(run_id, None) + self._run_to_graph_run.pop(run_id, None) + + task: dict[str, Any] = self._run_start_task.pop(run_id, {}) + task.update( + { + "ended_at": time.time(), + "status": "ERROR", + "stderr": str(error), + } + ) + if tel_start is not None: + task["telemetry_at_start"] = _tel_to_dict(tel_start) + self._interceptor.intercept_task(task) + self._record("llm_error", time.perf_counter() - t0) + + # --- Chat model events (same as LLM but different entry point) --- + + def on_chat_model_start( + self, + serialized: dict | None, + messages: list, + *, + run_id: UUID, + parent_run_id: UUID | None = None, + **kwargs: Any, + ) -> None: + """Buffer an llm_call task skeleton for a starting chat-model run.""" + t0 = time.perf_counter() + self._run_start[run_id] = time.time() + self._run_tel_start[run_id] = self._tel() + task_id = self._task_id_for(run_id) + # Propagate enclosing graph run mapping for group_id lookup in on_llm_end + if parent_run_id is not None: + graph_run_id = self._run_to_graph_run.get(parent_run_id) or ( + parent_run_id if parent_run_id in self._graph_runs else None + ) + if graph_run_id is not None: + self._run_to_graph_run[run_id] = graph_run_id + group_id = self._group_id_for(run_id) + parent_task_id = self._run_to_task.get(parent_run_id) if parent_run_id else None + model_name = self._model_name(serialized) + # Serialize messages: each message is a list of BaseMessage objects + serialized_messages = _safe_clip( + [[m.content if hasattr(m, "content") else str(m) for m in turn] for turn in messages] + ) + skeleton: dict[str, Any] = { + "task_id": task_id, + "subtype": "llm_call", + "activity_id": model_name, + "group_id": group_id, + "started_at": self._run_start[run_id], + "used": {"messages": serialized_messages, "model": model_name}, + "custom_metadata": {"model": model_name}, + } + if parent_task_id: + skeleton["parent_task_id"] = parent_task_id + self._run_start_task[run_id] = skeleton + self._record("chat_model_start", time.perf_counter() - t0) + + # on_chat_model_end fires on_llm_end in practice for most LangChain models + + # --- Tool events --- + + def on_tool_start( + self, + serialized: dict | None, + input_str: str, + *, + run_id: UUID, + parent_run_id: UUID | None = None, + **kwargs: Any, + ) -> None: + """Buffer a tool_call task skeleton for a starting tool run.""" + t0 = time.perf_counter() + self._run_start[run_id] = time.time() + self._run_tel_start[run_id] = self._tel() + tool_name = (serialized or {}).get("name") or "unknown_tool" + self._task_id_for(run_id) + parent_task_id = self._run_to_task.get(parent_run_id) if parent_run_id else None + # Propagate enclosing graph run mapping for group_id lookup in on_tool_end + if parent_run_id is not None: + graph_run_id = self._run_to_graph_run.get(parent_run_id) or ( + parent_run_id if parent_run_id in self._graph_runs else None + ) + if graph_run_id is not None: + self._run_to_graph_run[run_id] = graph_run_id + group_id = self._group_id_for(run_id) + task_id = self._run_to_task[run_id] + # Buffer skeleton — emitted as ONE complete record in on_tool_end + skeleton: dict[str, Any] = { + "task_id": task_id, + "subtype": "tool_call", + "activity_id": tool_name, + "group_id": group_id, + "started_at": self._run_start[run_id], + "used": {"input": _safe_clip(input_str)}, + "custom_metadata": {"tool_name": tool_name}, + } + if parent_task_id: + skeleton["parent_task_id"] = parent_task_id + self._run_start_task[run_id] = skeleton + self._record("tool_start", time.perf_counter() - t0) + + def on_tool_end( + self, + output: str, + *, + run_id: UUID, + parent_run_id: UUID | None = None, + **kwargs: Any, + ) -> None: + """Emit the completed tool_call task with its output and telemetry.""" + t0 = time.perf_counter() + tel_start = self._run_tel_start.pop(run_id, None) + tel_end = self._tel() + self._run_start.pop(run_id, None) + self._run_to_task.pop(run_id, None) + self._run_to_graph_run.pop(run_id, None) + + task: dict[str, Any] = self._run_start_task.pop(run_id, {}) + task.update( + { + "ended_at": time.time(), + "status": "FINISHED", + "generated": {"output": _safe_clip(output)}, + } + ) + if tel_start is not None: + task["telemetry_at_start"] = _tel_to_dict(tel_start) + if tel_end is not None: + task["telemetry_at_end"] = _tel_to_dict(tel_end) + + self._interceptor.intercept_task(task) + self._record("tool_end", time.perf_counter() - t0) + + def on_tool_error( + self, + error: BaseException, + *, + run_id: UUID, + parent_run_id: UUID | None = None, + **kwargs: Any, + ) -> None: + """Emit the failed tool_call task with ERROR status.""" + t0 = time.perf_counter() + tel_start = self._run_tel_start.pop(run_id, None) + self._run_start.pop(run_id, None) + self._run_to_task.pop(run_id, None) + self._run_to_graph_run.pop(run_id, None) + + task: dict[str, Any] = self._run_start_task.pop(run_id, {}) + task.update( + { + "ended_at": time.time(), + "status": "ERROR", + "stderr": str(error), + } + ) + if tel_start is not None: + task["telemetry_at_start"] = _tel_to_dict(tel_start) + self._interceptor.intercept_task(task) + self._record("tool_error", time.perf_counter() - t0) + + +def _build_handler_class(): + """Lazily build a concrete BaseCallbackHandler subclass. + + Keeps LangChain/LangGraph imported only when the plugin is actually used. + """ + from langchain_core.callbacks import BaseCallbackHandler + + class _ConcreteHandler(FlowceptLangGraphCallback, BaseCallbackHandler): + """Concrete LangChain callback handler with FlowCept provenance.""" + + def __init__(self, interceptor: LangGraphInterceptor, stats: _ProvenanceStats | None) -> None: + # BaseCallbackHandler.__init__ must be called before our __init__ + # because it sets ignore_* flags that LangChain checks. + BaseCallbackHandler.__init__(self) + FlowceptLangGraphCallback.__init__(self, interceptor, stats) + + # Tell LangChain we handle all event types + ignore_llm = False + ignore_chain = False + ignore_agent = False + ignore_chat_model = False + + return _ConcreteHandler + + +_HANDLER_CLASS = None + + +# --------------------------------------------------------------------------- +# Module-level active interceptor — set by start() / from_academy_plugin() +# --------------------------------------------------------------------------- + +_ACTIVE_INTERCEPTOR = None +_PROV_STATS: _ProvenanceStats | None = None + + +@contextmanager +def _timed(category: str): + t0 = time.perf_counter() + try: + yield + finally: + if _PROV_STATS is not None: + _PROV_STATS.record(category, time.perf_counter() - t0) + + +def record_llm_call(payload: dict) -> None: + """ + Public API to record an LLM call into the active FlowCept provenance graph. + + Converts the payload into a TaskObject (subtype=llm_call) and routes it + through the active interceptor. No-ops if the plugin has not been started. + + Minimum payload keys: + type : "chat_completion" + model : str + text : str + usage : dict with prompt_tokens / completion_tokens / total_tokens + """ + interceptor = _ACTIVE_INTERCEPTOR + if interceptor is None: + return + import uuid as _uuid + + with _timed("record_llm_call"): + elapsed = payload.get("elapsed_s", 0.0) + now = time.time() + model = payload.get("model_used") or payload.get("model", "unknown") + + used: dict = {} + for k in ( + "model", + "model_used", + "messages", + "user_prompt", + "system_prompt", + "temperature", + "top_p", + "max_tokens", + "reasoning_effort", + "tools_provided", + "tool_choice", + "stop_sequences", + "top_k", + "thinking_budget_tokens", + ): + if k in payload: + used[k] = payload[k] + if payload.get("temperature_suppressed"): + used["temperature_suppressed"] = True + + generated: dict = {} + for k in ( + "text", + "finish_reason", + "stop_reason", + "stop_sequence", + "tool_calls", + "tool_uses", + "thinking_text", + "system_fingerprint", + "response_id", + "usage", + "elapsed_s", + ): + if k in payload: + generated[k] = payload[k] + if "error" in payload: + generated["error"] = str(payload["error"]) + + task: dict = { + "task_id": str(_uuid.uuid4()), + "subtype": "llm_call", + "activity_id": model, + "started_at": now - elapsed, + "ended_at": now, + "status": "ERROR" if "error" in payload else "FINISHED", + "used": used, + "generated": generated, + "custom_metadata": { + "model": model, + "framework": payload.get("context", {}).get("framework", ""), + "context": payload.get("context", {}), + }, + } + interceptor.intercept_task(task) + + +def openai_chat( + prompt: str, + model: str = "gpt-4o-mini", + system: str = "You are a helpful assistant.", + temperature: float | None = 0.3, + top_p: float | None = None, + max_tokens: int | None = None, + n: int = 1, + stop: list[str] | str | None = None, + frequency_penalty: float = 0.0, + presence_penalty: float = 0.0, + seed: int | None = None, + reasoning_effort: str | None = None, + response_format: dict | None = None, + tools: list | None = None, + tool_choice: str | dict | None = None, + user: str | None = None, + context: dict | None = None, +) -> str: + """ + Make an OpenAI chat completion call and record it for FlowCept provenance. + + Captures all request parameters and response fields as a child TaskObject. + No-ops gracefully if OPENAI_API_KEY is not set or the plugin is not started. + """ + import openai as _openai + import os as _os + import time as _time + + client = _openai.OpenAI(api_key=_os.environ.get("OPENAI_API_KEY")) + messages = [ + {"role": "system", "content": system}, + {"role": "user", "content": prompt}, + ] + req: dict = {"model": model, "messages": messages, "n": n} + if temperature is not None: + req["temperature"] = temperature + if top_p is not None: + req["top_p"] = top_p + if max_tokens is not None: + req["max_completion_tokens"] = max_tokens + if stop is not None: + req["stop"] = stop + if frequency_penalty != 0.0: + req["frequency_penalty"] = frequency_penalty + if presence_penalty != 0.0: + req["presence_penalty"] = presence_penalty + if seed is not None: + req["seed"] = seed + if reasoning_effort is not None: + req["reasoning_effort"] = reasoning_effort + if response_format is not None: + req["response_format"] = response_format + if tools is not None: + req["tools"] = tools + if tool_choice is not None: + req["tool_choice"] = tool_choice + if user is not None: + req["user"] = user + + t0 = _time.time() + response = client.chat.completions.create(**req) + elapsed = _time.time() - t0 + + choice = response.choices[0] + text = choice.message.content or "" + usage = response.usage or {} + + usage_dict: dict = {} + for attr in ("prompt_tokens", "completion_tokens", "total_tokens"): + if hasattr(usage, attr): + usage_dict[attr] = getattr(usage, attr) + if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details: + ctd = usage.completion_tokens_details + usage_dict["reasoning_tokens"] = getattr(ctd, "reasoning_tokens", None) + usage_dict["accepted_prediction_tokens"] = getattr(ctd, "accepted_prediction_tokens", None) + usage_dict["rejected_prediction_tokens"] = getattr(ctd, "rejected_prediction_tokens", None) + if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details: + ptd = usage.prompt_tokens_details + usage_dict["cached_tokens"] = getattr(ptd, "cached_tokens", None) + + tool_calls = None + if choice.message.tool_calls: + tool_calls = [ + { + "id": tc.id, + "type": tc.type, + "function": {"name": tc.function.name, "arguments": tc.function.arguments}, + } + for tc in choice.message.tool_calls + ] + + record_llm_call( + { + "type": "chat_completion", + "model": model, + "model_used": response.model, + "messages": messages, + "user_prompt": prompt, + "system_prompt": system, + "temperature": temperature, + "temperature_suppressed": temperature is None, + "top_p": top_p, + "max_tokens": max_tokens, + "n": n, + "stop": stop, + "frequency_penalty": frequency_penalty, + "presence_penalty": presence_penalty, + "seed": seed, + "reasoning_effort": reasoning_effort, + "response_format": response_format, + "tools_provided": [t.get("function", {}).get("name") for t in (tools or [])], + "tool_choice": tool_choice, + "text": text, + "finish_reason": choice.finish_reason, + "tool_calls": tool_calls, + "system_fingerprint": getattr(response, "system_fingerprint", None), + "response_id": response.id, + "created": response.created, + "usage": usage_dict, + "elapsed_s": elapsed, + "context": context or {}, + } + ) + return text + + +def anthropic_chat( + prompt: str, + model: str = "claude-haiku-4-5-20251001", + system: str = "You are a helpful assistant.", + max_tokens: int = 1024, + temperature: float | None = 1.0, + top_p: float | None = None, + top_k: int | None = None, + stop_sequences: list[str] | None = None, + tools: list | None = None, + tool_choice: dict | None = None, + thinking: dict | None = None, + metadata: dict | None = None, + context: dict | None = None, +) -> str: + """ + Make an Anthropic (Claude) chat completion call and record it for FlowCept provenance. + + Captures all request/response fields — including thinking blocks, tool use, + and cache token counts — as a child TaskObject. + """ + import anthropic as _anthropic + import os as _os + import time as _time + + client = _anthropic.Anthropic(api_key=_os.environ.get("ANTHROPIC_API_KEY")) + req: dict = { + "model": model, + "max_tokens": max_tokens, + "system": system, + "messages": [{"role": "user", "content": prompt}], + } + if temperature is not None: + req["temperature"] = temperature + if top_p is not None: + req["top_p"] = top_p + if top_k is not None: + req["top_k"] = top_k + if stop_sequences: + req["stop_sequences"] = stop_sequences + if tools: + req["tools"] = tools + if tool_choice: + req["tool_choice"] = tool_choice + if thinking: + req["thinking"] = thinking + if metadata: + req["metadata"] = metadata + + t0 = _time.time() + response = client.messages.create(**req) + elapsed = _time.time() - t0 + + text = "" + thinking_text = "" + tool_uses = [] + for block in response.content: + if block.type == "text": + text += block.text + elif block.type == "thinking": + thinking_text += getattr(block, "thinking", "") + elif block.type == "tool_use": + tool_uses.append({"id": block.id, "name": block.name, "input": block.input}) + + usage = response.usage + usage_dict = { + "input_tokens": getattr(usage, "input_tokens", None), + "output_tokens": getattr(usage, "output_tokens", None), + "cache_creation_input_tokens": getattr(usage, "cache_creation_input_tokens", None), + "cache_read_input_tokens": getattr(usage, "cache_read_input_tokens", None), + } + + record_llm_call( + { + "type": "chat_completion", + "model": model, + "model_used": response.model, + "messages": req["messages"], + "user_prompt": prompt, + "system_prompt": system, + "max_tokens": max_tokens, + "temperature": temperature, + "temperature_suppressed": temperature is None, + "top_p": top_p, + "top_k": top_k, + "stop_sequences": stop_sequences, + "thinking_budget_tokens": (thinking or {}).get("budget_tokens"), + "tools_provided": [t.get("name") for t in (tools or [])], + "tool_choice": tool_choice, + "text": text, + "thinking_text": thinking_text if thinking_text else None, + "finish_reason": response.stop_reason, + "stop_sequence": response.stop_sequence, + "tool_uses": tool_uses if tool_uses else None, + "response_id": response.id, + "usage": usage_dict, + "elapsed_s": elapsed, + "context": context or {}, + } + ) + return text + + +# --------------------------------------------------------------------------- +# FlowceptAnthropicClient — wraps anthropic.Anthropic / AsyncAnthropic to +# capture full provenance for every messages.create / messages.stream call. +# --------------------------------------------------------------------------- + + +class FlowceptAnthropicClient: + """Anthropic client wrapper that records every model call as FlowCept provenance. + + Wraps an ``anthropic.Anthropic`` (or ``AsyncAnthropic``) client and records + every ``messages.create`` / ``messages.stream`` call as a FlowCept + provenance record (subtype=llm_call) via ``record_llm_call()``. + + Usage:: + + import anthropic + from flowcept.agents.langgraph.langgraph_plugin import FlowceptAnthropicClient + + client = FlowceptAnthropicClient(anthropic.Anthropic(), agent_name="my-agent") + response = client.messages.create( + model="claude-haiku-4-5-20251001", + max_tokens=1024, + messages=[{"role": "user", "content": "Hello"}], + ) + """ + + def __init__(self, client, agent_name=None, context=None): + self._inner = client + self._agent_name = agent_name + self._context: dict = context or {} + self.messages = _FlowceptAnthropicMessages(client.messages, agent_name, self._context) + + def __getattr__(self, name): + """Delegate every other attribute to the wrapped Anthropic client.""" + return getattr(self._inner, name) + + +class _FlowceptAnthropicMessages: + def __init__(self, messages_resource, agent_name, context): + self._inner = messages_resource + self._agent_name = agent_name + self._context = context + + def __getattr__(self, name): + return getattr(self._inner, name) + + def create(self, **kwargs): + import time as _time + + t0 = _time.time() + result = self._inner.create(**kwargs) + self._record(kwargs, result, _time.time() - t0) + return result + + async def async_create(self, **kwargs): + import time as _time + + t0 = _time.time() + result = await self._inner.create(**kwargs) + self._record(kwargs, result, _time.time() - t0) + return result + + def stream(self, **kwargs): + import time as _time + + return _FlowceptAnthropicStream(self._inner.stream(**kwargs), kwargs, self._record, _time.time()) + + def _record(self, kwargs, result, elapsed): + try: + model = kwargs.get("model", "unknown") + text = "" + thinking_text = "" + tool_uses = [] + for block in getattr(result, "content", []): + btype = getattr(block, "type", None) + if btype == "text": + text += getattr(block, "text", "") + elif btype == "thinking": + thinking_text += getattr(block, "thinking", "") + elif btype == "tool_use": + tool_uses.append( + { + "id": getattr(block, "id", None), + "name": getattr(block, "name", None), + "input": getattr(block, "input", None), + } + ) + usage = getattr(result, "usage", None) + usage_dict = ( + { + attr: getattr(usage, attr, None) + for attr in ( + "input_tokens", + "output_tokens", + "cache_creation_input_tokens", + "cache_read_input_tokens", + ) + } + if usage + else {} + ) + ctx = dict(self._context) + if self._agent_name: + ctx["agent_name"] = self._agent_name + payload = { + "type": "chat_completion", + "model": model, + "model_used": getattr(result, "model", model), + "messages": kwargs.get("messages"), + "system_prompt": kwargs.get("system"), + "max_tokens": kwargs.get("max_tokens"), + "temperature": kwargs.get("temperature"), + "top_p": kwargs.get("top_p"), + "top_k": kwargs.get("top_k"), + "stop_sequences": kwargs.get("stop_sequences"), + "thinking_budget_tokens": (kwargs.get("thinking") or {}).get("budget_tokens"), + "tools_provided": [t.get("name") for t in (kwargs.get("tools") or [])], + "tool_choice": kwargs.get("tool_choice"), + "text": text, + "thinking_text": thinking_text or None, + "finish_reason": getattr(result, "stop_reason", None), + "stop_sequence": getattr(result, "stop_sequence", None), + "tool_uses": tool_uses or None, + "response_id": getattr(result, "id", None), + "usage": usage_dict, + "elapsed_s": elapsed, + "context": ctx, + } + record_llm_call({k: v for k, v in payload.items() if v is not None}) + except Exception: + pass + + +class _FlowceptAnthropicStream: + def __init__(self, ctx_mgr, kwargs, record_fn, t0): + self._ctx_mgr = ctx_mgr + self._kwargs = kwargs + self._record_fn = record_fn + self._t0 = t0 + self._stream = None + + def __enter__(self): + self._stream = self._ctx_mgr.__enter__() + return self._stream + + def __exit__(self, *args): + import time as _time + + result = None + try: + result = self._stream.get_final_message() + except Exception: + pass + if result is not None: + self._record_fn(self._kwargs, result, _time.time() - self._t0) + return self._ctx_mgr.__exit__(*args) + + +# --------------------------------------------------------------------------- +# Public plugin class +# --------------------------------------------------------------------------- + + +class FlowceptLangGraphPlugin: + """ + FlowCept provenance plugin for LangGraph workflows. + + Captures the full provenance graph for each graph invocation: + graph run → node tasks → LLM/tool child tasks, with CPU/memory telemetry + and node enrichment. + + Parameters + ---------- + config : dict, optional + Plugin configuration keys: + enabled (bool, default True) + workflow_name (str, default "langgraph-workflow") + performance_tracking (bool, default True) + perf_csv (str, optional) — explicit path for timing CSV. + + Usage + ----- + Option A — explicit start/stop:: + + plugin = FlowceptLangGraphPlugin(config={...}) + plugin.start() + result = graph.invoke(state, config={"callbacks": [plugin.callback_handler]}) + plugin.stop() + + Option B — context manager:: + + with FlowceptLangGraphPlugin(config={...}) as plugin: + result = graph.invoke(state, config={"callbacks": [plugin.callback_handler]}) + """ + + def __init__(self, config: dict | None = None, _shared_interceptor=None) -> None: + cfg = config or {} + self._enabled: bool = cfg.get("enabled", True) + self._workflow_name: str = cfg.get("workflow_name", "langgraph-workflow") + self._campaign_id: str | None = cfg.get("campaign_id", None) + self._perf_tracking: bool = cfg.get("performance_tracking", True) + self._perf_csv: str | None = cfg.get("perf_csv", None) + # When _shared_interceptor is supplied the plugin does NOT start/stop + # its own Flowcept instance — it reuses the caller's buffer entirely. + self._shared_interceptor = _shared_interceptor + self._interceptor = _shared_interceptor or LangGraphInterceptor() + self._owns_interceptor: bool = _shared_interceptor is None + self._stats: _ProvenanceStats | None = None + self._handler: FlowceptLangGraphCallback | None = None + self._started = False + + @classmethod + def from_academy_plugin( + cls, + academy_plugin: Any, + config: dict | None = None, + ) -> "FlowceptLangGraphPlugin": + """Create a LangGraph plugin sharing an already-started FlowceptAcademyPlugin's buffer. + + All provenance records — from both Academy agents and LangGraph nodes — + are written to the exact same FlowCept in-memory buffer and end up in + the same JSONL file when the Academy plugin stops. + + Parameters + ---------- + academy_plugin : FlowceptAcademyPlugin + A plugin that has already been started (``plugin.start()`` called). + config : dict, optional + Same keys as FlowceptLangGraphPlugin.__init__ (enabled, + performance_tracking, perf_csv). + + Returns + ------- + FlowceptLangGraphPlugin + A started plugin whose ``callback_handler`` is ready for use. + + Example + ------- + :: + + academy_plugin = FlowceptAcademyPlugin(config={...}).start() + lg_plugin = FlowceptLangGraphPlugin.from_academy_plugin(academy_plugin) + # lg_plugin is already started — no need to call .start() + result = graph.invoke(state, config={"callbacks": [lg_plugin.callback_handler]}) + # Stop only the Academy plugin; it flushes the shared buffer. + academy_plugin.stop() + """ + global _ACTIVE_INTERCEPTOR, _HANDLER_CLASS, _PROV_STATS + # Inherit campaign_id from the academy plugin so all provenance is linked. + academy_interceptor = academy_plugin._interceptor + campaign_id = academy_interceptor._campaign_id + merged = dict(config or {}) + merged["campaign_id"] = campaign_id + # Create a proper LangGraphInterceptor (not an AcademyInterceptor). + lg_interceptor = LangGraphInterceptor() + lg_interceptor.start( + merged.get("workflow_name", "langgraph-workflow"), + campaign_id=campaign_id, + ) + inst = cls(config=merged, _shared_interceptor=lg_interceptor) + inst._interceptor = lg_interceptor + inst._owns_interceptor = True # we own this interceptor — stop it on stop() + inst._started = True + inst._stats = _ProvenanceStats() if merged.get("performance_tracking", True) else None + if _HANDLER_CLASS is None: + _HANDLER_CLASS = _build_handler_class() + inst._handler = _HANDLER_CLASS(lg_interceptor, inst._stats) + _ACTIVE_INTERCEPTOR = lg_interceptor + _PROV_STATS = inst._stats + return inst + + @property + def callback_handler(self) -> FlowceptLangGraphCallback: + """Return the LangChain callback handler to pass as ``config={"callbacks": [...]}``. + + Call ``start()`` before accessing this property. + """ + if self._handler is None: + raise RuntimeError("Plugin not started — call plugin.start() first.") + return self._handler + + def start(self) -> "FlowceptLangGraphPlugin": + """Start provenance capture and build the callback handler.""" + global _HANDLER_CLASS, _ACTIVE_INTERCEPTOR, _PROV_STATS + if not self._enabled or self._started: + return self + if not self._owns_interceptor: + # Already configured by from_academy_plugin(); nothing to start. + return self + try: + self._stats = _ProvenanceStats() if self._perf_tracking else None + _PROV_STATS = self._stats + self._interceptor.start(self._workflow_name, campaign_id=self._campaign_id) + self._campaign_id = self._interceptor._campaign_id + + if _HANDLER_CLASS is None: + _HANDLER_CLASS = _build_handler_class() + + self._handler = _HANDLER_CLASS(self._interceptor, self._stats) + self._started = True + _ACTIVE_INTERCEPTOR = self._interceptor + + wf_id = self._interceptor._workflow_id + campaign_id = self._interceptor._campaign_id + print( + f"[FlowceptLangGraphPlugin] Started\n" + f" workflow_id : {wf_id}\n" + f" campaign_id : {campaign_id}\n" + f" Capturing : graph runs (sub-workflows), node execution " + f"(parent_task_id), LLM + tool calls (child tasks, telemetry).", + flush=True, + ) + except Exception as e: + print( + f"[FlowceptLangGraphPlugin] WARNING: failed to start — {e!r}. Continuing without provenance capture.", + flush=True, + ) + _log.exception("FlowceptLangGraphPlugin start failed") + self._enabled = False + return self + + def stop(self) -> None: + """Stop provenance capture, flush the buffer, and report overhead stats.""" + global _ACTIVE_INTERCEPTOR, _PROV_STATS + if not self._started: + return + if not self._owns_interceptor: + # Buffer is owned by the Academy plugin — do not stop or flush it here. + # Just print the perf stats if enabled. + self._started = False + print( + "[FlowceptLangGraphPlugin] Detached from shared buffer (buffer flushed by the owning plugin).", + flush=True, + ) + self._maybe_write_perf_csv() + _ACTIVE_INTERCEPTOR = None + _PROV_STATS = None + return + try: + _t0 = time.perf_counter() + self._interceptor.stop() + if self._stats is not None: + self._stats.record("flush", time.perf_counter() - _t0) + except Exception as e: + print(f"[FlowceptLangGraphPlugin] Warning during stop: {e!r}", flush=True) + self._started = False + print("[FlowceptLangGraphPlugin] Stopped.", flush=True) + self._maybe_write_perf_csv() + _PROV_STATS = None + _ACTIVE_INTERCEPTOR = None + + def _maybe_write_perf_csv(self) -> None: + + if self._stats is not None: + print( + "\n[FlowceptLangGraphPlugin] Provenance overhead report:\n" + + self._stats.summary() + + "\n (N = event count; Total/Mean/Min/Max in ms/µs respectively)\n", + flush=True, + ) + wf_id = self._interceptor._workflow_id + csv_path = self._perf_csv or f"langgraph_provenance_perf_{wf_id}.csv" + try: + self._stats.to_csv(csv_path, workflow_id=wf_id) + print( + f"[FlowceptLangGraphPlugin] Performance stats written to {csv_path}", + flush=True, + ) + except Exception as e: + print( + f"[FlowceptLangGraphPlugin] Warning: could not write perf CSV — {e!r}", + flush=True, + ) + + def __enter__(self) -> "FlowceptLangGraphPlugin": + """Start the plugin when entering a ``with`` block.""" + return self.start() + + def __exit__(self, *_: Any) -> None: + """Stop the plugin when leaving a ``with`` block.""" + self.stop() + + +# --------------------------------------------------------------------------- +# Serialisation helper (identical to flowcept_plugin.py) +# --------------------------------------------------------------------------- + + +def _safe_clip(obj: Any, _depth: int = 0) -> Any: + """Recursively convert objects to JSON-serialisable form without truncation.""" + if _depth > 8: + return str(obj) + if obj is None or isinstance(obj, (bool, int, float, str)): + return obj + if isinstance(obj, dict): + return {str(k): _safe_clip(v, _depth + 1) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return [_safe_clip(v, _depth + 1) for v in obj] + return repr(obj) + + +def _tel_to_dict(tel: Any) -> dict | None: + if tel is None: + return None + try: + return tel.to_dict() + except Exception: + return None diff --git a/src/flowcept/agents/mcp/mcp_server.py b/src/flowcept/agents/mcp/mcp_server.py index 5f5b8545..45aaffe5 100644 --- a/src/flowcept/agents/mcp/mcp_server.py +++ b/src/flowcept/agents/mcp/mcp_server.py @@ -11,6 +11,7 @@ # Import all mcp_tools modules so their @mcp_flowcept.tool() decorators fire from flowcept.agents.mcp.mcp_tools.session_tools import check_liveness +import flowcept.agents.mcp.mcp_tools.analysis_mcp_tools # noqa: F401 import flowcept.agents.mcp.mcp_tools.db_query_mcp_tools # noqa: F401 import flowcept.agents.mcp.mcp_tools.dashboard_mcp_tools # noqa: F401 import flowcept.agents.mcp.mcp_tools.df_query_mcp_tools # noqa: F401 diff --git a/src/flowcept/agents/mcp/mcp_tools/analysis_mcp_tools.py b/src/flowcept/agents/mcp/mcp_tools/analysis_mcp_tools.py new file mode 100644 index 00000000..c6637210 --- /dev/null +++ b/src/flowcept/agents/mcp/mcp_tools/analysis_mcp_tools.py @@ -0,0 +1,186 @@ +"""Thin MCP wrappers for provenance analysis tools. + +All analysis logic lives in :mod:`flowcept.agents.prov_analysis`. The ``df_*`` +tools run over the records loaded in the agent's in-memory context (the same +context ``df_query_mcp_tools`` queries); the ``db_*`` variants pull records +from the database via ``DBAPI``. ``compare_executions`` compares two +workflows, preferring in-memory records and falling back to the DB. +""" + +from typing import Any, Dict, List, Optional + +from flowcept.agents.mcp.context_manager import EMPTY_DF_MESSAGE, ctx_manager, get_df_context, mcp_flowcept +from flowcept.agents.prov_analysis import tools as _core +from flowcept.agents.tool_result import ToolResult +from flowcept.commons.vocabulary import PROV_AGENT +from flowcept.instrumentation.flowcept_agent_task import agent_flowcept_task + + +def _context_records() -> List[Dict[str, Any]]: + """Collect provenance records from the agent's in-memory context. + + Prefers the raw task dicts kept by the context manager (they retain nested + ``used``/``generated``/``custom_metadata`` fields); falls back to the + flattened tasks DataFrame when no raw tasks are held. + """ + records: List[Dict[str, Any]] = [] + workflow = ctx_manager.context.workflow_msg_obj + if workflow: + records.append(workflow) + tasks = ctx_manager.context.tasks or [] + if tasks: + records.extend(tasks) + else: + df, _, _, _ = get_df_context(context_kind="tasks") + if df is not None and len(df): + records.extend(df.to_dict("records")) + return records + + +def _db_records(workflow_id: Optional[str] = None) -> List[Dict[str, Any]]: + """Load workflow and task records for one workflow (or all tasks) from the DB.""" + from flowcept.flowcept_api.db_api import DBAPI + + db = DBAPI() + filter = {"workflow_id": workflow_id} if workflow_id else {} + records: List[Dict[str, Any]] = [] + for wf in db.workflow_query(filter=filter) or []: + wf = dict(wf) + wf.setdefault("type", "workflow") + records.append(wf) + for task in db.task_query(filter=filter) or []: + task = dict(task) + task.setdefault("type", "task") + records.append(task) + return records + + +def _db_analysis(tool_name: str, analysis_fn, workflow_id: Optional[str], **kwargs) -> ToolResult: + """Run one analysis over DB records, converting DB failures to a ToolResult error.""" + try: + records = _db_records(workflow_id) + except Exception as e: + return ToolResult(code=499, result=f"Error in {tool_name}: could not query DB: {e}", tool_name=tool_name) + return analysis_fn(records, **kwargs) + + +# --------------------------------------------------------------------------- +# DF variants — analyze the records loaded in the agent's in-memory context +# --------------------------------------------------------------------------- + + +@mcp_flowcept.tool() +@agent_flowcept_task(subtype=PROV_AGENT.AGENT_TOOL) +def df_summarize_execution(workflow_id: str = None) -> ToolResult: + """Summarize the loaded execution: counts by activity/subtype, statuses, durations, token usage. + + Optionally pass ``workflow_id`` to restrict the summary to one workflow. + """ + records = _context_records() + if not records: + return ToolResult(code=404, result=EMPTY_DF_MESSAGE, tool_name="df_summarize_execution") + return _core.summarize_execution(records, workflow_id=workflow_id) + + +@mcp_flowcept.tool() +@agent_flowcept_task(subtype=PROV_AGENT.AGENT_TOOL) +def df_analyze_errors() -> ToolResult: + """Analyze failures in the loaded execution: per-activity error rates and stderr excerpts.""" + records = _context_records() + if not records: + return ToolResult(code=404, result=EMPTY_DF_MESSAGE, tool_name="df_analyze_errors") + return _core.analyze_errors(records) + + +@mcp_flowcept.tool() +@agent_flowcept_task(subtype=PROV_AGENT.AGENT_TOOL) +def df_agent_behavior() -> ToolResult: + """Profile per-agent behavior in the loaded execution: turns, tool calls, LLM calls, token usage.""" + records = _context_records() + if not records: + return ToolResult(code=404, result=EMPTY_DF_MESSAGE, tool_name="df_agent_behavior") + return _core.analyze_agent_behavior(records) + + +@mcp_flowcept.tool() +@agent_flowcept_task(subtype=PROV_AGENT.AGENT_TOOL) +def df_find_slowest(limit: int = 10) -> ToolResult: + """Find the slowest tasks of the loaded execution, longest elapsed first.""" + records = _context_records() + if not records: + return ToolResult(code=404, result=EMPTY_DF_MESSAGE, tool_name="df_find_slowest") + return _core.find_slowest_tasks(records, limit=limit) + + +@mcp_flowcept.tool() +@agent_flowcept_task(subtype=PROV_AGENT.AGENT_TOOL) +def df_cross_framework_links() -> ToolResult: + """List cross-framework provenance links (source_agent_id edges) in the loaded execution.""" + records = _context_records() + if not records: + return ToolResult(code=404, result=EMPTY_DF_MESSAGE, tool_name="df_cross_framework_links") + return _core.cross_framework_links(records) + + +# --------------------------------------------------------------------------- +# DB variants — analyze records pulled from the database via DBAPI +# --------------------------------------------------------------------------- + + +@mcp_flowcept.tool() +@agent_flowcept_task(subtype=PROV_AGENT.AGENT_TOOL) +def db_summarize_execution(workflow_id: str = None) -> ToolResult: + """Summarize an execution from DB records: counts, statuses, durations, token usage.""" + return _db_analysis("db_summarize_execution", _core.summarize_execution, workflow_id, workflow_id=workflow_id) + + +@mcp_flowcept.tool() +@agent_flowcept_task(subtype=PROV_AGENT.AGENT_TOOL) +def db_analyze_errors(workflow_id: str = None) -> ToolResult: + """Analyze failures from DB records: per-activity error rates and stderr excerpts.""" + return _db_analysis("db_analyze_errors", _core.analyze_errors, workflow_id) + + +@mcp_flowcept.tool() +@agent_flowcept_task(subtype=PROV_AGENT.AGENT_TOOL) +def db_agent_behavior(workflow_id: str = None) -> ToolResult: + """Profile per-agent behavior from DB records: turns, tool calls, LLM calls, token usage.""" + return _db_analysis("db_agent_behavior", _core.analyze_agent_behavior, workflow_id) + + +@mcp_flowcept.tool() +@agent_flowcept_task(subtype=PROV_AGENT.AGENT_TOOL) +def db_find_slowest(workflow_id: str = None, limit: int = 10) -> ToolResult: + """Find the slowest tasks from DB records, longest elapsed first.""" + return _db_analysis("db_find_slowest", _core.find_slowest_tasks, workflow_id, limit=limit) + + +@mcp_flowcept.tool() +@agent_flowcept_task(subtype=PROV_AGENT.AGENT_TOOL) +def db_cross_framework_links(workflow_id: str = None) -> ToolResult: + """List cross-framework provenance links (source_agent_id edges) from DB records.""" + return _db_analysis("db_cross_framework_links", _core.cross_framework_links, workflow_id) + + +# --------------------------------------------------------------------------- +# Comparison — two workflow executions +# --------------------------------------------------------------------------- + + +@mcp_flowcept.tool() +@agent_flowcept_task(subtype=PROV_AGENT.AGENT_TOOL) +def compare_executions(workflow_id_a: str, workflow_id_b: str) -> ToolResult: + """Compare two workflow executions per activity: count, duration, and error-rate deltas. + + Records for each workflow are taken from the agent's in-memory context when + present there, otherwise queried from the database. + """ + in_memory = _context_records() + + def records_for(workflow_id: str) -> List[Dict[str, Any]]: + matching = [r for r in in_memory if r.get("workflow_id") == workflow_id] + if matching: + return matching + return _db_records(workflow_id) + + return _core.compare_executions(records_for(workflow_id_a), records_for(workflow_id_b)) diff --git a/src/flowcept/agents/openai_agents/openai_agents_plugin.py b/src/flowcept/agents/openai_agents/openai_agents_plugin.py new file mode 100644 index 00000000..9030eca9 --- /dev/null +++ b/src/flowcept/agents/openai_agents/openai_agents_plugin.py @@ -0,0 +1,303 @@ +"""OpenAI Agents SDK wrapper. + +The Agents SDK already traces itself — every run produces a trace of typed +spans — and lets you register extra processors for those spans. So capture is a +processor, not a wrapper: nothing about how you call the SDK changes. + + from flowcept.agents.openai_agents.openai_agents_plugin import install + install() + + result = await Runner.run(agent, "fix the failing test") + +How the SDK's span types map onto PROV-AGENT: + +=================== ==================================================== +trace the session workflow, plus one turn spanning the run +``generation`` ``ai_model_invocation`` at call granularity +``response`` the same, for the Responses API +``function`` ``agent_tool`` +``mcp_tools`` ``agent_tool`` +``guardrail`` ``agent_tool`` (it runs, it passes or trips) +nested ``agent`` a subagent workflow — agent-as-tool +root ``agent`` the session itself; contributes its name, not a record +``handoff`` a lifecycle event on the session +=================== ==================================================== + +Duck-typed against the SDK: nothing here imports ``agents``, so the module +imports and tests without it. +""" + +from __future__ import annotations + +import datetime +from typing import Any + +from flowcept.agents.harness.config import Config +from flowcept.agents.harness.tracer import SessionTracer + +#: Span types recorded as tool executions. +TOOL_SPANS = frozenset({"function", "mcp_tools", "guardrail", "custom"}) + +#: Span types recorded as model invocations. +MODEL_SPANS = frozenset({"generation", "response"}) + + +class FlowceptTraceProcessor: + """An Agents SDK ``TracingProcessor`` that writes Flowcept provenance. + + Deliberately not subclassing the SDK's ``TracingProcessor``: doing so would + make ``openai-agents`` a hard import, and the SDK duck-types processors. + """ + + def __init__(self, config: Config | None = None, *, harness: str = "openai_agents"): + self.config = config + self.harness = harness + self._tracers: dict[str, SessionTracer] = {} + #: span_id -> trace_id, so a span can find its session. + self._traces: dict[str, str] = {} + self._parents: dict[str, str | None] = {} + #: span_id of agent spans that opened a subagent workflow. + self._agents: dict[str, str] = {} + + # -- traces -------------------------------------------------------------- + + def on_trace_start(self, trace: Any) -> None: + """Open a session tracer for a new SDK trace.""" + trace_id = getattr(trace, "trace_id", None) + if not trace_id: + return + # `group_id` is the SDK's conversation/thread id when the caller set + # one. Preferring it means a multi-turn conversation is one session + # rather than one session per run. + session_id = getattr(trace, "group_id", None) or trace_id + tracer = SessionTracer(self.harness, str(session_id), config=self.config) + self._tracers[trace_id] = tracer + tracer.start(source=getattr(trace, "name", None) or "run") + tracer.prompt(None, prompt_id=trace_id) + + def on_trace_end(self, trace: Any) -> None: + """Close the trace's session tracer and drop its span bookkeeping.""" + trace_id = getattr(trace, "trace_id", None) + tracer = self._tracers.pop(trace_id, None) + if tracer is None: + return + tracer.turn_end() + # A grouped conversation gets another trace, and the session is reopened + # by it; closing here is still right, because the closing record + # supersedes rather than duplicates. + tracer.end() + # Drop the trace's spans. Computed before the loop, since the first + # mapping cleared is the one the list is derived from. + span_ids = [k for k, v in self._traces.items() if v == trace_id] + for mapping in (self._traces, self._parents, self._agents): + for span_id in span_ids: + mapping.pop(span_id, None) + + # -- spans --------------------------------------------------------------- + + def on_span_start(self, span: Any) -> None: + """Record the start of an SDK span in the owning session tracer.""" + tracer, data, span_type = self._resolve(span) + if tracer is None: + return + + span_id = getattr(span, "span_id", "") + self._traces[span_id] = getattr(span, "trace_id", "") + self._parents[span_id] = getattr(span, "parent_id", None) + + if span_type == "agent" and getattr(span, "parent_id", None): + self._agents[span_id] = tracer.subagent_start( + getattr(data, "name", None) or "agent", + agent_ref=span_id, + ) + elif span_type == "agent": + # The root agent is the session; name it rather than nest it. + model = getattr(data, "model", None) + if isinstance(model, str): + tracer.model = model + elif span_type in TOOL_SPANS: + tracer.tool_start( + _span_name(data, span_type), + _as_json(getattr(data, "input", None) or getattr(data, "data", None)), + tool_use_id=span_id, + agent_ref=self._owning_agent(span), + ) + + def on_span_end(self, span: Any) -> None: + """Record the completion of an SDK span in the owning session tracer.""" + tracer, data, span_type = self._resolve(span) + if tracer is None: + return + + span_id = getattr(span, "span_id", "") + error = _error_text(getattr(span, "error", None)) + started = _parse_time(getattr(span, "started_at", None)) + # The span already measured its own duration; without both ends the + # record would be stamped with the time the callback happened to run. + ended = _parse_time(getattr(span, "ended_at", None)) + when = {"timestamp": ended} if ended else {} + + if span_type == "agent": + ref = self._agents.pop(span_id, None) + if ref: + tracer.subagent_stop(ref, response=_as_text(getattr(data, "output", None)), error=error, **when) + elif span_type in TOOL_SPANS: + tracer.tool_end( + span_id, + name=_span_name(data, span_type), + tool_response=_as_json(getattr(data, "output", None)) if error is None else None, + error=error or _guardrail_error(data), + agent_ref=self._owning_agent(span), + started_at=started, + **when, + ) + elif span_type in MODEL_SPANS: + model, usage, output = _model_details(data) + tracer.llm_call( + model=model, + prompt=_as_text(getattr(data, "input", None)), + response=output, + usage=usage, + call_id=span_id, + started_at=started, + error=error, + **when, + ) + elif span_type == "handoff": + source = getattr(data, "from_agent", None) + target = getattr(data, "to_agent", None) + tracer.notify(f"handoff: {source} -> {target}", source="handoff") + + self._traces.pop(span_id, None) + self._parents.pop(span_id, None) + + # -- plumbing ------------------------------------------------------------ + + def _resolve(self, span: Any) -> tuple[SessionTracer | None, Any, str | None]: + tracer = self._tracers.get(getattr(span, "trace_id", None)) + data = getattr(span, "span_data", None) + return tracer, data, getattr(data, "type", None) + + def _owning_agent(self, span: Any) -> str | None: + """Return the nearest enclosing subagent, so its tools land in its workflow.""" + parent = getattr(span, "parent_id", None) + seen = 0 + while parent and seen < 32: # depth guard: the chain comes from outside + if parent in self._agents: + return self._agents[parent] + parent = self._parents.get(parent) + seen += 1 + return None + + def shutdown(self) -> None: + """Close any session tracers still open when the SDK shuts down.""" + for trace_id in list(self._tracers): + tracer = self._tracers.pop(trace_id) + tracer.end(source="shutdown") + + def force_flush(self) -> None: + """Do nothing; records are emitted as spans end.""" + return None + + +def install(config: Config | None = None, *, replace: bool = False) -> FlowceptTraceProcessor: + """Register the processor with the Agents SDK's tracing provider. + + Adds to the existing processors by default, so the SDK's own trace export + keeps working. Pass ``replace=True`` to make Flowcept the only consumer. + """ + from agents import tracing + + processor = FlowceptTraceProcessor(config) + if replace: + tracing.set_trace_processors([processor]) + else: + tracing.add_trace_processor(processor) + return processor + + +# -- span-data readers -------------------------------------------------------- + + +def _span_name(data: Any, span_type: str | None) -> str: + name = getattr(data, "name", None) + if isinstance(name, str) and name: + return name + return span_type or "tool" + + +def _guardrail_error(data: Any) -> str | None: + """Return an error message for a tripped guardrail, which is a failed tool.""" + if getattr(data, "type", None) == "guardrail" and getattr(data, "triggered", False): + return "guardrail triggered" + return None + + +def _model_details(data: Any) -> tuple[str | None, dict[str, Any] | None, str | None]: + """Pull model, usage, and output text from a generation or response span.""" + model = getattr(data, "model", None) + usage = _as_dict(getattr(data, "usage", None)) + output = getattr(data, "output", None) + + response = getattr(data, "response", None) + if response is not None: + model = model or getattr(response, "model", None) + usage = usage or _as_dict(getattr(response, "usage", None)) + output = output if output is not None else getattr(response, "output_text", None) + + return (model if isinstance(model, str) else None), usage, _as_text(output) + + +def _error_text(error: Any) -> str | None: + if error is None: + return None + if isinstance(error, dict): + message = error.get("message") + data = error.get("data") + return f"{message}: {data}" if message and data else (message or _as_text(data)) + return _as_text(error) + + +def _parse_time(value: Any) -> float | None: + """Span timestamps are ISO-8601 strings.""" + if isinstance(value, (int, float)): + return float(value) + if not isinstance(value, str): + return None + try: + return datetime.datetime.fromisoformat(value.replace("Z", "+00:00")).timestamp() + except ValueError: + return None + + +def _as_dict(value: Any) -> dict[str, Any] | None: + if value is None or isinstance(value, dict): + return value + for attribute in ("model_dump", "to_dict", "_asdict"): + method = getattr(value, attribute, None) + if callable(method): + try: + result = method() + except Exception: + continue + if isinstance(result, dict): + return result + return getattr(value, "__dict__", None) or None + + +def _as_json(value: Any) -> Any: + """Keep structure where there is any; ``used`` reads better with fields.""" + if value is None or isinstance(value, (dict, list, str, int, float, bool)): + return value + return _as_dict(value) or repr(value) + + +def _as_text(value: Any) -> str | None: + import json + + if value is None or isinstance(value, str): + return value + try: + return json.dumps(value, default=repr) + except (TypeError, ValueError): + return repr(value) diff --git a/src/flowcept/agents/otel/otel_plugin.py b/src/flowcept/agents/otel/otel_plugin.py new file mode 100644 index 00000000..d12aeead --- /dev/null +++ b/src/flowcept/agents/otel/otel_plugin.py @@ -0,0 +1,361 @@ +"""OpenTelemetry ingest. + +Many harnesses and agent frameworks already emit OTel spans and have no hook +system at all. Rather than ask them to add one, this adapter reads spans and +turns them into the same provenance everything else produces. + +Two ways in: + +*A span exporter* (:class:`FlowceptSpanExporter`) plugs into an in-process +tracer provider, so anything instrumented with OTel starts producing Flowcept +provenance with three lines of setup. + +*A file reader* (:func:`ingest_file`) consumes spans already written as JSON, +which is what ``OTEL_TRACES_EXPORTER=console`` and most collectors produce. + +The mapping follows the OTel GenAI semantic conventions, which name the +attributes this cares about: ``gen_ai.operation.name`` distinguishes a model +call from a tool call, ``gen_ai.tool.name`` names the tool, and +``gen_ai.conversation.id`` groups spans into a session. Non-GenAI spans are +ignored -- an HTTP client span is not provenance. + +``gen_ai.system`` is optional and never affects grouping: session identity +derives from the conversation id alone, so spans with and without the +attribute land in one workflow. The first non-empty value a conversation +shows is recorded once as a lifecycle event; later or conflicting values are +ignored (first-wins). +""" + +from __future__ import annotations + +import json +from collections.abc import Iterable +from pathlib import Path +from typing import Any + +from flowcept.agents.harness.config import Config, load_config +from flowcept.agents.harness.events import HarnessEvent +from flowcept.agents.harness.recorder import Recorder +from flowcept.agents.harness.vocab import EventKind + +#: Session grouping, most specific first. +SESSION_KEYS = ( + "gen_ai.conversation.id", + "gen_ai.session.id", + "session.id", + "session_id", + "thread.id", +) + +#: `gen_ai.operation.name` values that mean "the model was invoked". +MODEL_OPERATIONS = frozenset({"chat", "generate_content", "text_completion", "embeddings", "generate"}) + +#: ...and the ones that mean "a tool ran". +TOOL_OPERATIONS = frozenset({"execute_tool", "invoke_tool", "tool"}) + +#: The first non-empty ``gen_ai.system`` seen per conversation. The provider +#: name must never feed workflow identity (mixed presence would split one +#: conversation across workflows), so it is tracked here and recorded once as +#: a lifecycle event instead. First non-empty value wins. +_session_systems: dict[str, str] = {} + + +def _attr(attributes: dict[str, Any], *keys: str) -> Any: + for key in keys: + value = attributes.get(key) + if value is not None: + return value + return None + + +def _nanos_to_seconds(value: Any) -> float | None: + """OTel timestamps are nanoseconds since the epoch.""" + if isinstance(value, (int, float)): + return value / 1e9 + if isinstance(value, str): + # Console exporters sometimes emit ISO-8601 instead. + try: + import datetime + + return datetime.datetime.fromisoformat(value.replace("Z", "+00:00")).timestamp() + except ValueError: + return None + return None + + +def span_to_event(span: dict[str, Any], *, harness: str = "otel") -> HarnessEvent | None: + """Map one span dict onto a normalized event, or ``None`` to skip it. + + Accepts both the shape produced by the SDK's console exporter and the + flattened shape most collectors emit. + """ + attributes = span.get("attributes") or {} + if not isinstance(attributes, dict): + return None + + session_id = _attr(attributes, *SESSION_KEYS) + if not session_id: + # Without a conversation id spans cannot be grouped into a session, and + # a per-span workflow would be noise rather than provenance. + return None + + operation = _attr(attributes, "gen_ai.operation.name", "gen_ai.operation", "operation.name") + tool_name = _attr(attributes, "gen_ai.tool.name", "tool.name") + + if operation in TOOL_OPERATIONS or tool_name: + kind = EventKind.TOOL_POST + elif operation in MODEL_OPERATIONS: + kind = EventKind.LLM_CALL + else: + return None + + status = span.get("status") or {} + status_code = status.get("status_code") if isinstance(status, dict) else status + error = None + if str(status_code).upper() in ("ERROR", "STATUS_CODE_ERROR"): + error = (status.get("description") if isinstance(status, dict) else None) or "span reported an error" + if kind == EventKind.TOOL_POST: + kind = EventKind.TOOL_ERROR + + usage = { + key: attributes[full] + for key, full in ( + ("input_tokens", "gen_ai.usage.input_tokens"), + ("output_tokens", "gen_ai.usage.output_tokens"), + ("total_tokens", "gen_ai.usage.total_tokens"), + ) + if full in attributes + } + + started = _nanos_to_seconds(span.get("start_time") or span.get("startTimeUnixNano")) + ended = _nanos_to_seconds(span.get("end_time") or span.get("endTimeUnixNano")) + + return HarnessEvent( + kind=kind, + # Deliberately NOT `gen_ai.system`: the harness partitions workflow + # identity, and the provider attribute may be set on only some of a + # conversation's spans. See `_provider_notice`. + harness=harness, + session_id=str(session_id), + # The event carries the span's *end*; `started_at` preserves the + # duration the span already measured. + timestamp=ended or started or None, + started_at=started, + model=_attr(attributes, "gen_ai.request.model", "gen_ai.response.model", "llm.model_name"), + prompt=_as_text(_attr(attributes, "gen_ai.prompt", "gen_ai.input.messages", "input.value")), + response=_as_text(_attr(attributes, "gen_ai.completion", "gen_ai.output.messages", "output.value")), + tool_name=tool_name or span.get("name"), + tool_use_id=_attr(attributes, "gen_ai.tool.call.id", "tool.call.id") or _span_id(span), + tool_input=_as_json(_attr(attributes, "gen_ai.tool.call.arguments", "tool.arguments", "input.value")), + tool_response=_as_json(_attr(attributes, "gen_ai.tool.call.result", "tool.result", "output.value")), + error=error, + call_id=_attr(attributes, "gen_ai.response.id") or _span_id(span), + usage=usage or None, + agent_name=_attr(attributes, "gen_ai.agent.name", "agent.name"), + agent_ref=_attr(attributes, "gen_ai.agent.id", "agent.id"), + ) + + +def _provider_notice(event: HarnessEvent, attributes: dict[str, Any]) -> HarnessEvent | None: + """Build a one-time lifecycle event recording the session's ``gen_ai.system``. + + The provider name must not partition the session (that would split spans + with and without the attribute across workflows), so it lands as a + ``harness_event`` task in the conversation's workflow instead. The first + non-empty value wins; later or conflicting values return ``None``. + """ + system = _attr(attributes, "gen_ai.system", "service.name") + if not system or event.session_id in _session_systems: + return None + _session_systems[event.session_id] = str(system) + return HarnessEvent( + kind=EventKind.NOTIFICATION, + harness=event.harness, + session_id=event.session_id, + timestamp=event.timestamp, + source="gen_ai.system", + message=str(system), + raw={"gen_ai.system": str(system)}, + ) + + +def _span_id(span: dict[str, Any]) -> str | None: + context = span.get("context") or span.get("spanContext") or {} + if isinstance(context, dict): + value = context.get("span_id") or context.get("spanId") + if value: + return str(value) + value = span.get("span_id") or span.get("spanId") + return str(value) if value else None + + +def _as_text(value: Any) -> str | None: + if value is None or isinstance(value, str): + return value + return json.dumps(value, default=repr) + + +def _as_json(value: Any) -> Any: + """Parse attribute values that carry JSON as a string. + + OTel attributes are scalars, so structured tool arguments arrive + JSON-encoded. Recovering the structure means ``used`` holds real fields + rather than one opaque string. + """ + if isinstance(value, str): + stripped = value.strip() + if stripped[:1] in ("{", "["): + try: + return json.loads(stripped) + except ValueError: + return value + return value + + +def ingest_spans(spans: Iterable[dict[str, Any]], config: Config | None = None) -> int: + """Record every GenAI span in ``spans``. Returns how many were recorded.""" + config = config or load_config() + recorder = Recorder(config) + recorded = 0 + for span in spans: + if not isinstance(span, dict): + continue + event = span_to_event(span) + if event is None: + continue + if recorder.record(event): + recorded += 1 + notice = _provider_notice(event, span.get("attributes") or {}) + if notice is not None: + recorder.record(notice) + return recorded + + +def ingest_file(path: str | Path, config: Config | None = None) -> int: + """Record spans from a JSON or JSONL file of exported spans.""" + path = Path(path) + text = path.read_text(encoding="utf-8") + + spans: list[dict[str, Any]] = [] + stripped = text.lstrip() + if stripped.startswith("["): + loaded = json.loads(text) + spans = [s for s in loaded if isinstance(s, dict)] + else: + # JSONL, or the console exporter's stream of pretty-printed objects. + decoder = json.JSONDecoder() + index = 0 + while index < len(text): + while index < len(text) and text[index] in " \t\r\n": + index += 1 + if index >= len(text): + break + try: + obj, index = decoder.raw_decode(text, index) + except ValueError: + break + if isinstance(obj, dict): + spans.append(obj) + + return ingest_spans(_flatten(spans), config) + + +def _flatten(spans: Iterable[dict[str, Any]]) -> Iterable[dict[str, Any]]: + """Yield individual spans from either bare spans or OTLP envelopes.""" + for item in spans: + if "resourceSpans" in item or "resource_spans" in item: + envelopes = item.get("resourceSpans") or item.get("resource_spans") or [] + for resource in envelopes: + for scope in resource.get("scopeSpans") or resource.get("scope_spans") or []: + for span in scope.get("spans") or []: + yield _normalize_otlp(span) + else: + yield item + + +def _normalize_otlp(span: dict[str, Any]) -> dict[str, Any]: + """Convert OTLP's list-of-key-value attributes into a plain dict.""" + attributes = span.get("attributes") + if isinstance(attributes, list): + flat: dict[str, Any] = {} + for entry in attributes: + key = entry.get("key") + value = entry.get("value") + if key is None or not isinstance(value, dict): + continue + # OTLP wraps each value in a type tag: {"stringValue": "..."}. + for tag in ("stringValue", "intValue", "doubleValue", "boolValue"): + if tag in value: + flat[key] = value[tag] + break + span = {**span, "attributes": flat} + return span + + +class FlowceptSpanExporter: + """An OTel ``SpanExporter`` that records spans as Flowcept provenance. + + Usage:: + + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from flowcept.agents.otel.otel_plugin import FlowceptSpanExporter + + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(FlowceptSpanExporter())) + + Deliberately not subclassing ``SpanExporter``: doing so would make + ``opentelemetry-sdk`` a hard import of this module, and the class is + duck-typed by the SDK anyway. + """ + + def __init__(self, config: Config | None = None): + self.config = config or load_config() + self._recorder = Recorder(self.config) + + def export(self, spans) -> Any: + """Convert each readable span to an event and record it.""" + for span in spans: + try: + data = self._readable_to_dict(span) + event = span_to_event(data) + except Exception: + continue + if event is not None: + self._recorder.record(event) + notice = _provider_notice(event, data.get("attributes") or {}) + if notice is not None: + self._recorder.record(notice) + try: + from opentelemetry.sdk.trace.export import SpanExportResult + + return SpanExportResult.SUCCESS + except ImportError: + return None + + @staticmethod + def _readable_to_dict(span) -> dict[str, Any]: + """Adapt a ``ReadableSpan`` to the dict shape :func:`span_to_event` reads.""" + context = span.get_span_context() if hasattr(span, "get_span_context") else None + status = getattr(span, "status", None) + return { + "name": getattr(span, "name", None), + "attributes": dict(getattr(span, "attributes", None) or {}), + "start_time": getattr(span, "start_time", None), + "end_time": getattr(span, "end_time", None), + "context": {"span_id": format(context.span_id, "016x")} if context else {}, + "status": { + "status_code": getattr(getattr(status, "status_code", None), "name", None), + "description": getattr(status, "description", None), + } + if status is not None + else {}, + } + + def shutdown(self) -> None: + """Do nothing; the recorder needs no teardown.""" + return None + + def force_flush(self, timeout_millis: int = 30000) -> bool: + """Report success; records are written as spans are exported.""" + return True diff --git a/src/flowcept/agents/prov_analysis/README.md b/src/flowcept/agents/prov_analysis/README.md new file mode 100644 index 00000000..dad8aa21 --- /dev/null +++ b/src/flowcept/agents/prov_analysis/README.md @@ -0,0 +1,51 @@ +# `prov_analysis/` + +Agentic provenance analysis over Flowcept PROV-AGENT records. Pure-Python analysis functions shared by **four surfaces**: the harness MCP server (`agents/harness/mcp_server.py`), the Flowcept agent MCP server (`mcp/mcp_tools/analysis_mcp_tools.py`), the webservice chat (`chat_orchestration/tool_registry.py`), and the `flowcept-harness analyze` CLI subcommand. + +## Why this exists + +Every surface that can see provenance records — harness buffers, the agent's in-memory context, or the DB — needs the same analyses (what happened, what failed, what was slow, how agents behaved, how frameworks link). Putting the logic here once means all surfaces stay consistent, following the same layering rule as `data_query_tools/`: **cores are framework-free; MCP/LangChain wrappers are thin**. + +## Modules + +### `core.py` +Pure functions over `records: list[dict]` — a mixed list of workflow/task/agent records in either the harness-buffer shape (`agents/harness/prov.py`) or the framework-plugin shape (LangGraph, CrewAI, AutoGen, Academy). Dependency-light: stdlib plus `flowcept.report.aggregations` (itself stdlib-only), so the harness MCP server can import it lazily without pulling pandas or a backend. Fields that may be absent degrade gracefully. + +- `load_records(jsonl_path=None, records=None, workflow_id=None, campaign_id=None)` — load from a JSONL buffer (via `report.loaders.read_jsonl`), pass records through, or (guarded, lazy) load from the Flowcept DB. +- `summarize_execution(records, workflow_id=None)` — counts by type/subtype/activity, duration bounds, status counts, campaigns/agents seen, token-usage totals (harness `custom_metadata.llm_usage` and plugin `generated.*_tokens` fields), and per-activity rows via `report.aggregations.group_activities`. +- `analyze_errors(records)` — failed tasks grouped by `activity_id` with stderr/message excerpts, error rate per activity, first/last failure times. +- `analyze_agent_behavior(records)` — per agent/session: turns, tool calls by tool, LLM calls, token usage, error counts, avg/max task durations; plus session workflows and subagent counts. +- `find_slowest_tasks(records, limit=10)` — slowest tasks with `task_id`, `activity_id`, `elapsed_seconds`, `status`, and parent-chain depth. +- `cross_framework_links(records)` — edges built from `source_agent_id` pointers (top-level, `custom_metadata.source_agent_id`, or `used`/`used.inputs._source_agent_id`): `{source_task_id, target_task_id, target_workflow_id, frameworks}`, plus the count of unlinked tasks. +- `compare_executions(records_a, records_b)` — per-activity count/duration/error-rate deltas between two executions. + +### `tools.py` +`ToolResult` wrappers over each core function, following the `_guarded` convention of `data_query_tools` (3xx success dicts, 4xx error strings). Framework-free: no MCP or LangChain imports. + +## Surfaces + +``` +harness MCP server (stdlib-only file; imports core lazily inside tools) + analyze_session / analyze_errors / find_slowest / cross_links + └─► prov_analysis.core over the session's JSONL buffer + +Flowcept agent MCP (mcp/mcp_tools/analysis_mcp_tools.py) + df_summarize_execution, df_analyze_errors, df_agent_behavior, + df_find_slowest, df_cross_framework_links ── agent in-memory context + db_* variants ── DBAPI workflow/task queries + compare_executions(workflow_id_a, workflow_id_b) + └─► prov_analysis.tools ─► prov_analysis.core + +chat (chat_orchestration/tool_registry.py) + StructuredTool wrappers routing df/db by tool_context, via run_mcp + +CLI + flowcept-harness analyze [--errors | --slowest N | --links] + └─► prov_analysis.core over the session's JSONL buffer +``` + +## Record-shape notes + +- Harness records carry `type: task|workflow|agent`, subtypes `ai_model_invocation` / `agent_tool` / `harness_event`, and token usage under `custom_metadata.llm_usage`. +- Framework-plugin records may omit `type` (they always carry `task_id`), use subtypes like `langgraph_graph|langgraph_node|llm_call|tool_call`, and put token usage in `generated.prompt_tokens|completion_tokens|total_tokens`. +- Cross-framework links travel as `custom_metadata.source_agent_id` (and the raw `_source_agent_id` inside `used.inputs`). diff --git a/src/flowcept/agents/prov_analysis/__init__.py b/src/flowcept/agents/prov_analysis/__init__.py new file mode 100644 index 00000000..d577ed04 --- /dev/null +++ b/src/flowcept/agents/prov_analysis/__init__.py @@ -0,0 +1,27 @@ +"""Agentic provenance analysis over Flowcept PROV-AGENT records. + +Pure, dependency-light analysis functions (:mod:`.core`) plus ToolResult +wrappers (:mod:`.tools`). Exposed through the harness MCP server, the +Flowcept agent MCP server, the LangChain chat tool registry, and the +``flowcept-harness analyze`` CLI subcommand. +""" + +from flowcept.agents.prov_analysis.core import ( + analyze_agent_behavior, + analyze_errors, + compare_executions, + cross_framework_links, + find_slowest_tasks, + load_records, + summarize_execution, +) + +__all__ = [ + "load_records", + "summarize_execution", + "analyze_errors", + "analyze_agent_behavior", + "find_slowest_tasks", + "cross_framework_links", + "compare_executions", +] diff --git a/src/flowcept/agents/prov_analysis/core.py b/src/flowcept/agents/prov_analysis/core.py new file mode 100644 index 00000000..931f1bd9 --- /dev/null +++ b/src/flowcept/agents/prov_analysis/core.py @@ -0,0 +1,565 @@ +"""Pure analysis functions over PROV-AGENT provenance records. + +Every function here takes ``records: list[dict]`` — a mixed list of workflow, +task, and agent records as produced either by the harness buffer +(``flowcept.agents.harness``) or by the framework plugins (LangGraph, CrewAI, +AutoGen, Academy). Records are inspected per record via their ``type`` / +``subtype`` fields, and any field that may be absent degrades gracefully. + +This module is dependency-light on purpose: stdlib plus +:mod:`flowcept.report.aggregations` (itself stdlib-only), so the harness MCP +server can import it lazily without pulling pandas or any backend. +""" + +from __future__ import annotations + +from collections import Counter, defaultdict +from typing import Any, Dict, List, Optional + +from flowcept.report.aggregations import ( + as_float, + elapsed_seconds, + fmt_timestamp_utc, + group_activities, + workflow_bounds, +) + +__all__ = [ + "load_records", + "summarize_execution", + "analyze_errors", + "analyze_agent_behavior", + "find_slowest_tasks", + "cross_framework_links", + "compare_executions", +] + +_STATUS_ERROR = "ERROR" + +#: Token-usage keys observed in real records (harness ``custom_metadata.llm_usage`` +#: and framework-plugin ``generated`` fields). +_USAGE_KEYS = ( + "prompt_tokens", + "completion_tokens", + "total_tokens", + "input_tokens", + "output_tokens", + "reasoning_tokens", +) + + +# -- record classification ----------------------------------------------------- + + +def _is_workflow(record: Dict[str, Any]) -> bool: + return record.get("type") == "workflow" + + +def _is_agent(record: Dict[str, Any]) -> bool: + return record.get("type") == "agent" + + +def _is_task(record: Dict[str, Any]) -> bool: + # Harness records always carry ``type: task``; framework-plugin task dicts + # may omit ``type`` but always carry a ``task_id``. + if record.get("type") == "task": + return True + return record.get("type") is None and record.get("task_id") is not None + + +def _tasks(records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + return [r for r in records if _is_task(r)] + + +def _workflows(records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + return [r for r in records if _is_workflow(r)] + + +def _filter_workflow(records: List[Dict[str, Any]], workflow_id: Optional[str]) -> List[Dict[str, Any]]: + if not workflow_id: + return records + return [r for r in records if r.get("workflow_id") == workflow_id] + + +def _custom_metadata(record: Dict[str, Any]) -> Dict[str, Any]: + meta = record.get("custom_metadata") + return meta if isinstance(meta, dict) else {} + + +def _task_usage(record: Dict[str, Any]) -> Dict[str, float]: + """Extract token-usage numbers from one task record. + + Looks in ``custom_metadata.llm_usage`` (harness shape) and in ``generated`` + (framework-plugin shape, e.g. ``generated.total_tokens``). + """ + usage: Dict[str, float] = {} + sources: List[Dict[str, Any]] = [] + llm_usage = _custom_metadata(record).get("llm_usage") + if isinstance(llm_usage, dict): + sources.append(llm_usage) + generated = record.get("generated") + if isinstance(generated, dict): + sources.append(generated) + for source in sources: + for key in _USAGE_KEYS: + val = as_float(source.get(key)) + if val is not None: + usage[key] = usage.get(key, 0.0) + val + return usage + + +def _sum_usage(tasks: List[Dict[str, Any]]) -> Dict[str, Any]: + totals: Dict[str, float] = {} + tasks_with_usage = 0 + for task in tasks: + usage = _task_usage(task) + if not usage: + continue + tasks_with_usage += 1 + for key, val in usage.items(): + totals[key] = totals.get(key, 0.0) + val + return {"totals": {k: int(v) for k, v in totals.items()}, "n_tasks_with_usage": tasks_with_usage} + + +def _framework_of(record: Dict[str, Any]) -> Optional[str]: + """Best-effort name of the framework/harness a record came from.""" + harness = _custom_metadata(record).get("harness") + if harness: + return str(harness) + adapter = record.get("adapter_id") + if adapter: + return str(adapter).rsplit(".", 1)[-1] + subtype = record.get("subtype") + if isinstance(subtype, str) and "_" in subtype and subtype.split("_", 1)[0] in ("langgraph", "crewai", "autogen"): + return subtype.split("_", 1)[0] + return None + + +# -- loading ------------------------------------------------------------------- + + +def load_records( + jsonl_path: Optional[str] = None, + records: Optional[List[Dict[str, Any]]] = None, + workflow_id: Optional[str] = None, + campaign_id: Optional[str] = None, +) -> List[Dict[str, Any]]: + """Load provenance records from a JSONL buffer, a list, or the Flowcept DB. + + Parameters + ---------- + jsonl_path : str, optional + Path to a JSONL buffer file (harness buffer or dumped Flowcept buffer). + records : list of dict, optional + Already-loaded records; returned as-is (list-copied). + workflow_id : str, optional + When neither ``jsonl_path`` nor ``records`` is given, load this + workflow from the Flowcept DB (requires a reachable backend). + campaign_id : str, optional + Same as ``workflow_id`` but for a whole campaign. + + Returns + ------- + list of dict + Mixed workflow/task/agent records. + """ + if records is not None: + return list(records) + if jsonl_path is not None: + from pathlib import Path + + from flowcept.report.loaders import read_jsonl + + loaded, _skipped = read_jsonl(Path(jsonl_path)) + return loaded + if workflow_id or campaign_id: + # Guarded, lazy DB loading: only reached when explicitly requested. + from flowcept.report.loaders import load_records_from_db + + dataset = load_records_from_db(workflow_id=workflow_id, campaign_id=campaign_id) + out: List[Dict[str, Any]] = [] + for wf in dataset.get("workflows") or ([dataset["workflow"]] if dataset.get("workflow") else []): + wf = dict(wf) + wf.setdefault("type", "workflow") + out.append(wf) + for task in dataset.get("tasks") or []: + task = dict(task) + task.setdefault("type", "task") + out.append(task) + return out + return [] + + +# -- analyses ------------------------------------------------------------------ + + +def summarize_execution(records: List[Dict[str, Any]], workflow_id: Optional[str] = None) -> Dict[str, Any]: + """Summarize an execution: counts, duration bounds, statuses, agents, token usage. + + Parameters + ---------- + records : list of dict + Mixed workflow/task records (harness or framework-plugin shape). + workflow_id : str, optional + Restrict the summary to records of one workflow. + + Returns + ------- + dict + Counts by type/subtype/activity, duration bounds, status counts, + campaigns and agents seen, and token-usage totals where present. + """ + records = _filter_workflow(records, workflow_id) + tasks = _tasks(records) + workflows = _workflows(records) + + by_type = Counter(str(r.get("type") or ("task" if _is_task(r) else "unknown")) for r in records) + by_subtype = Counter(str(t.get("subtype", "unknown")) for t in tasks) + by_activity = Counter(str(t.get("activity_id", "unknown")) for t in tasks) + status_counts = Counter(str(t.get("status", "unknown")) for t in tasks) + + min_start, max_end, total_elapsed = workflow_bounds(tasks) + if total_elapsed is None and workflows: + min_start, max_end, total_elapsed = workflow_bounds(workflows) + + campaigns = sorted({str(r["campaign_id"]) for r in records if r.get("campaign_id")}) + agent_ids = {str(r["agent_id"]) for r in records if r.get("agent_id")} + agent_names = sorted({str(r["name"]) for r in records if _is_agent(r) and r.get("name")}) + + return { + "n_records": len(records), + "n_workflows": len(workflows), + "n_tasks": len(tasks), + "counts_by_type": dict(by_type), + "tasks_by_subtype": dict(by_subtype), + "tasks_by_activity": dict(by_activity), + "status_counts": dict(status_counts), + "started_at": min_start, + "ended_at": max_end, + "started_at_utc": fmt_timestamp_utc(min_start) if min_start is not None else None, + "ended_at_utc": fmt_timestamp_utc(max_end) if max_end is not None else None, + "total_elapsed_seconds": total_elapsed, + "campaigns": campaigns, + "agents": sorted(agent_ids), + "agent_names": agent_names, + "token_usage": _sum_usage(tasks), + "activities": group_activities(tasks) if tasks else [], + } + + +def analyze_errors(records: List[Dict[str, Any]]) -> Dict[str, Any]: + """Analyze failed tasks: per-activity failures, error rates, first/last failure. + + Parameters + ---------- + records : list of dict + Mixed provenance records. + + Returns + ------- + dict + ``by_activity`` maps activity_id to failure count, error rate, and + stderr/message excerpts; plus overall counts and failure time bounds. + """ + tasks = _tasks(records) + failed = [t for t in tasks if str(t.get("status")) == _STATUS_ERROR] + + totals_by_activity = Counter(str(t.get("activity_id", "unknown")) for t in tasks) + by_activity: Dict[str, Dict[str, Any]] = {} + for task in failed: + activity = str(task.get("activity_id", "unknown")) + entry = by_activity.setdefault(activity, {"n_failed": 0, "error_rate": 0.0, "excerpts": []}) + entry["n_failed"] += 1 + message = task.get("stderr") or _custom_metadata(task).get("error") or task.get("stdout") + if message and len(entry["excerpts"]) < 5: + entry["excerpts"].append(str(message)[:300]) + for activity, entry in by_activity.items(): + total = totals_by_activity.get(activity, 0) + entry["n_total"] = total + entry["error_rate"] = round(entry["n_failed"] / total, 4) if total else None + + failure_starts = [as_float(t.get("started_at")) for t in failed] + failure_starts = [s for s in failure_starts if s is not None] + first = min(failure_starts) if failure_starts else None + last = max(failure_starts) if failure_starts else None + + return { + "n_tasks": len(tasks), + "n_failed": len(failed), + "overall_error_rate": round(len(failed) / len(tasks), 4) if tasks else None, + "by_activity": by_activity, + "first_failure_at": first, + "last_failure_at": last, + "first_failure_at_utc": fmt_timestamp_utc(first) if first is not None else None, + "last_failure_at_utc": fmt_timestamp_utc(last) if last is not None else None, + } + + +def analyze_agent_behavior(records: List[Dict[str, Any]]) -> Dict[str, Any]: + """Profile agent behavior per agent/session. + + For each agent (falling back to the session workflow when no ``agent_id`` + is present): turns, tool calls by tool, LLM calls, token usage, and + avg/max task durations. Session workflows and subagent counts are + reported alongside. + + Parameters + ---------- + records : list of dict + Mixed provenance records. + + Returns + ------- + dict + ``agents`` keyed by agent/session id, plus ``sessions`` and + ``n_subagent_sessions``. + """ + tasks = _tasks(records) + workflows = _workflows(records) + + agents: Dict[str, Dict[str, Any]] = {} + durations: Dict[str, List[float]] = defaultdict(list) + for task in tasks: + key = str(task.get("agent_id") or task.get("workflow_id") or "unknown") + entry = agents.setdefault( + key, + { + "turns": 0, + "llm_calls": 0, + "tool_calls": 0, + "tool_calls_by_tool": {}, + "token_usage": {}, + "n_tasks": 0, + "n_errors": 0, + }, + ) + entry["n_tasks"] += 1 + if str(task.get("status")) == _STATUS_ERROR: + entry["n_errors"] += 1 + subtype = str(task.get("subtype", "")) + granularity = _custom_metadata(task).get("granularity") + if subtype == "ai_model_invocation": + if granularity == "call": + entry["llm_calls"] += 1 + else: + entry["turns"] += 1 + elif subtype == "llm_call": + entry["llm_calls"] += 1 + elif subtype in ("agent_tool", "tool_call"): + entry["tool_calls"] += 1 + tool = str(task.get("activity_id", "unknown")) + entry["tool_calls_by_tool"][tool] = entry["tool_calls_by_tool"].get(tool, 0) + 1 + for usage_key, val in _task_usage(task).items(): + entry["token_usage"][usage_key] = int(entry["token_usage"].get(usage_key, 0) + val) + elapsed = elapsed_seconds(task.get("started_at"), task.get("ended_at")) + if elapsed is not None: + durations[key].append(elapsed) + + for key, entry in agents.items(): + vals = durations.get(key) or [] + entry["avg_task_seconds"] = round(sum(vals) / len(vals), 4) if vals else None + entry["max_task_seconds"] = round(max(vals), 4) if vals else None + + subagent_workflows = [w for w in workflows if w.get("parent_workflow_id")] + subagents_by_parent = Counter(str(w.get("parent_workflow_id")) for w in subagent_workflows) + sessions = [] + for wf in workflows: + if wf.get("parent_workflow_id"): + continue + wid = str(wf.get("workflow_id", "unknown")) + sessions.append( + { + "workflow_id": wid, + "name": wf.get("name"), + "subtype": wf.get("subtype"), + "status": wf.get("status"), + "elapsed_seconds": elapsed_seconds(wf.get("started_at"), wf.get("ended_at")), + "n_subagents": subagents_by_parent.get(wid, 0), + } + ) + + return { + "agents": agents, + "sessions": sessions, + "n_subagent_sessions": len(subagent_workflows), + } + + +def find_slowest_tasks(records: List[Dict[str, Any]], limit: int = 10) -> List[Dict[str, Any]]: + """Return the slowest tasks, longest elapsed first. + + Parameters + ---------- + records : list of dict + Mixed provenance records. + limit : int, optional + Maximum number of rows returned (default 10). + + Returns + ------- + list of dict + Rows with ``task_id``, ``activity_id``, ``subtype``, ``elapsed_seconds``, + ``status``, and ``parent_depth`` (length of the parent_task_id chain). + """ + tasks = _tasks(records) + by_id = {t.get("task_id"): t for t in tasks if t.get("task_id")} + + def depth(task: Dict[str, Any]) -> int: + seen = set() + d = 0 + current = task + while True: + parent_id = current.get("parent_task_id") + if not parent_id or parent_id in seen: + return d + seen.add(parent_id) + d += 1 + parent = by_id.get(parent_id) + if parent is None: + return d + current = parent + + rows = [] + for task in tasks: + elapsed = elapsed_seconds(task.get("started_at"), task.get("ended_at")) + if elapsed is None: + continue + rows.append( + { + "task_id": task.get("task_id"), + "activity_id": task.get("activity_id"), + "subtype": task.get("subtype"), + "elapsed_seconds": round(elapsed, 4), + "status": task.get("status"), + "parent_depth": depth(task), + } + ) + rows.sort(key=lambda r: r["elapsed_seconds"], reverse=True) + return rows[: max(0, int(limit))] + + +def cross_framework_links(records: List[Dict[str, Any]]) -> Dict[str, Any]: + """Extract cross-framework provenance links between records. + + A link exists when a task carries a source pointer — top-level + ``source_agent_id``, ``custom_metadata.source_agent_id``, or the raw + ``_source_agent_id`` key inside ``used`` / ``used.inputs`` — naming a task + or agent from another framework. + + Parameters + ---------- + records : list of dict + Mixed provenance records. + + Returns + ------- + dict + ``links`` (list of ``{source_task_id, target_task_id, + target_workflow_id, frameworks}``), ``n_links``, ``n_unlinked_tasks``, + and ``frameworks_seen``. + """ + tasks = _tasks(records) + by_id: Dict[str, Dict[str, Any]] = {} + for record in records: + for id_key in ("task_id", "agent_id", "workflow_id"): + rid = record.get(id_key) + if rid and rid not in by_id: + by_id[str(rid)] = record + + links = [] + linked_task_ids = set() + for task in tasks: + used = task.get("used") if isinstance(task.get("used"), dict) else {} + inputs = used.get("inputs") if isinstance(used.get("inputs"), dict) else {} + source = ( + task.get("source_agent_id") + or _custom_metadata(task).get("source_agent_id") + or used.get("_source_agent_id") + or inputs.get("_source_agent_id") + ) + if not source: + continue + source = str(source) + source_record = by_id.get(source) + frameworks = [fw for fw in (_framework_of(source_record) if source_record else None, _framework_of(task)) if fw] + links.append( + { + "source_task_id": source, + "target_task_id": task.get("task_id"), + "target_workflow_id": task.get("workflow_id"), + "frameworks": frameworks, + } + ) + linked_task_ids.add(task.get("task_id")) + + frameworks_seen = sorted({fw for r in records if (fw := _framework_of(r))}) + return { + "links": links, + "n_links": len(links), + "n_unlinked_tasks": len([t for t in tasks if t.get("task_id") not in linked_task_ids]), + "frameworks_seen": frameworks_seen, + } + + +def compare_executions(records_a: List[Dict[str, Any]], records_b: List[Dict[str, Any]]) -> Dict[str, Any]: + """Compare two executions per activity: count, duration, and error-rate deltas. + + Parameters + ---------- + records_a : list of dict + Records of the first execution. + records_b : list of dict + Records of the second execution. + + Returns + ------- + dict + ``activities`` keyed by activity_id with ``*_a``/``*_b``/``*_delta`` + columns, plus ``only_in_a``/``only_in_b`` and total bounds. + """ + tasks_a, tasks_b = _tasks(records_a), _tasks(records_b) + rows_a = {r["activity_id"]: r for r in group_activities(tasks_a)} + rows_b = {r["activity_id"]: r for r in group_activities(tasks_b)} + + def error_rate(row: Dict[str, Any]) -> Optional[float]: + n = row.get("n_tasks") or 0 + if not n: + return None + errors = (row.get("status_counts") or {}).get(_STATUS_ERROR, 0) + return round(errors / n, 4) + + activities: Dict[str, Dict[str, Any]] = {} + for activity in sorted(set(rows_a) | set(rows_b)): + a, b = rows_a.get(activity), rows_b.get(activity) + count_a = a["n_tasks"] if a else 0 + count_b = b["n_tasks"] if b else 0 + avg_a = a.get("elapsed_avg") if a else None + avg_b = b.get("elapsed_avg") if b else None + rate_a = error_rate(a) if a else None + rate_b = error_rate(b) if b else None + activities[activity] = { + "count_a": count_a, + "count_b": count_b, + "count_delta": count_b - count_a, + "elapsed_avg_a": avg_a, + "elapsed_avg_b": avg_b, + "elapsed_avg_delta": (avg_b - avg_a) if avg_a is not None and avg_b is not None else None, + "error_rate_a": rate_a, + "error_rate_b": rate_b, + "error_rate_delta": (rate_b - rate_a) if rate_a is not None and rate_b is not None else None, + } + + _, _, total_a = workflow_bounds(tasks_a) + _, _, total_b = workflow_bounds(tasks_b) + return { + "activities": activities, + "only_in_a": sorted(set(rows_a) - set(rows_b)), + "only_in_b": sorted(set(rows_b) - set(rows_a)), + "totals": { + "n_tasks_a": len(tasks_a), + "n_tasks_b": len(tasks_b), + "total_elapsed_a": total_a, + "total_elapsed_b": total_b, + "total_elapsed_delta": (total_b - total_a) if total_a is not None and total_b is not None else None, + }, + } diff --git a/src/flowcept/agents/prov_analysis/tools.py b/src/flowcept/agents/prov_analysis/tools.py new file mode 100644 index 00000000..969b4390 --- /dev/null +++ b/src/flowcept/agents/prov_analysis/tools.py @@ -0,0 +1,79 @@ +"""ToolResult wrappers over :mod:`flowcept.agents.prov_analysis.core`. + +Framework-free (no MCP, no LangChain imports): each function wraps one core +analysis in the ``ToolResult`` convention used by ``data_query_tools``, so both +the MCP surface and the chat tool registry can reuse them without drift. +""" + +from __future__ import annotations + +from typing import Any, Dict, List, Optional + +from flowcept.agents.prov_analysis import core +from flowcept.agents.tool_result import ToolResult + + +def _guarded(tool_name: str): + """Decorate a tool function: convert exceptions to ToolResult error codes.""" + + def decorator(func): + def wrapper(*args, **kwargs): + try: + return func(*args, **kwargs) + except ValueError as e: + return ToolResult(code=400, result=str(e), tool_name=tool_name) + except Exception as e: + from flowcept.commons.flowcept_logger import FlowceptLogger + + FlowceptLogger().exception(e) + return ToolResult(code=499, result=f"Error in {tool_name}: {e}", tool_name=tool_name) + + wrapper.__name__ = func.__name__ + wrapper.__doc__ = func.__doc__ + return wrapper + + return decorator + + +@_guarded("summarize_execution") +def summarize_execution(records: List[Dict[str, Any]], workflow_id: Optional[str] = None) -> ToolResult: + """Summarize an execution (counts, durations, statuses, agents, token usage).""" + return ToolResult( + code=301, + result=core.summarize_execution(records, workflow_id=workflow_id), + tool_name="summarize_execution", + ) + + +@_guarded("analyze_errors") +def analyze_errors(records: List[Dict[str, Any]]) -> ToolResult: + """Analyze failed tasks: per-activity failure counts, error rates, excerpts.""" + return ToolResult(code=301, result=core.analyze_errors(records), tool_name="analyze_errors") + + +@_guarded("analyze_agent_behavior") +def analyze_agent_behavior(records: List[Dict[str, Any]]) -> ToolResult: + """Profile per-agent behavior: turns, tool calls, LLM calls, token usage.""" + return ToolResult(code=301, result=core.analyze_agent_behavior(records), tool_name="analyze_agent_behavior") + + +@_guarded("find_slowest_tasks") +def find_slowest_tasks(records: List[Dict[str, Any]], limit: int = 10) -> ToolResult: + """Return the slowest tasks, longest elapsed first.""" + return ToolResult( + code=301, + result={"tasks": core.find_slowest_tasks(records, limit=limit)}, + tool_name="find_slowest_tasks", + ) + + +@_guarded("cross_framework_links") +def cross_framework_links(records: List[Dict[str, Any]]) -> ToolResult: + """List cross-framework provenance links (source_agent_id edges).""" + return ToolResult(code=301, result=core.cross_framework_links(records), tool_name="cross_framework_links") + + +@_guarded("compare_executions") +def compare_executions(records_a: List[Dict[str, Any]], records_b: List[Dict[str, Any]]) -> ToolResult: + """Compare two executions per activity (count/duration/error-rate deltas).""" + return ToolResult(code=301, result=core.compare_executions(records_a, records_b), tool_name="compare_executions") diff --git a/src/flowcept/commons/autoflush_buffer.py b/src/flowcept/commons/autoflush_buffer.py index c1dfb3c2..0bca1a13 100644 --- a/src/flowcept/commons/autoflush_buffer.py +++ b/src/flowcept/commons/autoflush_buffer.py @@ -23,10 +23,10 @@ def __init__( self._swap_event = Event() self._stop_event = Event() - self._timer_thread = Thread(target=self.time_based_flush) + self._timer_thread = Thread(target=self.time_based_flush, daemon=True) self._timer_thread.start() - self._flush_thread = Thread(target=self._flush_buffers) + self._flush_thread = Thread(target=self._flush_buffers, daemon=True) self._flush_thread.start() self._flush_function = flush_function diff --git a/src/flowcept/commons/daos/mq_dao/mq_dao_base.py b/src/flowcept/commons/daos/mq_dao/mq_dao_base.py index 44159beb..b4da0a5b 100644 --- a/src/flowcept/commons/daos/mq_dao/mq_dao_base.py +++ b/src/flowcept/commons/daos/mq_dao/mq_dao_base.py @@ -54,6 +54,10 @@ def build(*args, **kwargs) -> "MQDao": from flowcept.commons.daos.mq_dao.mq_dao_mofka import MQDaoMofka return MQDaoMofka(*args, **kwargs) + elif MQ_TYPE == "diaspora": + from flowcept.commons.daos.mq_dao.mq_dao_diaspora import MQDaoDiaspora + + return MQDaoDiaspora(*args, **kwargs) elif MQ_TYPE == "rabbitmq": from flowcept.commons.daos.mq_dao.mq_dao_rabbitmq import MQDaoRabbitMQ diff --git a/src/flowcept/commons/daos/mq_dao/mq_dao_diaspora.py b/src/flowcept/commons/daos/mq_dao/mq_dao_diaspora.py new file mode 100644 index 00000000..6b4ffdbd --- /dev/null +++ b/src/flowcept/commons/daos/mq_dao/mq_dao_diaspora.py @@ -0,0 +1,102 @@ +import uuid +from typing import Callable + +import msgpack +from time import time + +from diaspora_stream.api import Driver, Validator + +from flowcept.commons.daos.mq_dao.mq_dao_base import MQDao +from flowcept.configs import MQ_SETTINGS, MQ_CHANNEL + + +class MQDaoDiaspora(MQDao): + """Main class to communicate with diaspora.""" + + driver_options = { + "root_path": "/tmp/diaspora-data/", + } + _validator = Validator.from_metadata() + _driver = Driver(backend="files", options=driver_options) + _topic = _driver.open_topic(MQ_SETTINGS["channel"]) + + def __init__(self, adapter_settings=None, with_producer=True): + super().__init__(adapter_settings=adapter_settings) + self.producer = None + if with_producer: + producer_name = "p" + MQ_CHANNEL + "-" + str(uuid.uuid4())[:8] + print(f"Starting producer {producer_name}") + self.producer = MQDaoDiaspora._topic.producer(producer_name) + + def subscribe(self): + """Subscribe to Diaspora topic.""" + self.consumer = MQDaoDiaspora._topic.consumer(name=MQ_CHANNEL + str(uuid.uuid4())) + + def message_listener(self, message_handler: Callable): + """Diaspora's Message listener.""" + try: + while True: + future = self.consumer.pull() + event = future.wait(timeout_ms=1) + while not future.completed: + event = future.wait(timeout_ms=1) + message = event.metadata + self.logger.debug(f"Received message: {message}") + if not message_handler(message): + break + except Exception as e: + self.logger.exception(e) + finally: + pass + + def send_message(self, message: dict, channel=MQ_CHANNEL, serializer=msgpack.dumps): + """Send a single message to Diaspora.""" + self.producer.push(metadata=message) # using metadata to send data + self.producer.flush() + + def _send_message_timed(self, message: dict, channel=MQ_CHANNEL, serializer=msgpack.dumps): + t1 = time() + self.send_message(message, channel, serializer) + t2 = time() + self._flush_events.append(["single", t1, t2, t2 - t1, len(str(message).encode())]) + + def _bulk_publish(self, buffer, channel=MQ_CHANNEL, serializer=msgpack.dumps): + try: + for m in buffer: + self.producer.push(metadata=m) + except Exception as e: + self.logger.exception(e) + self.logger.error("Some messages couldn't be flushed! Check the messages' contents!") + self.logger.error(f"Message that caused error: {buffer}") + try: + self.producer.flush() + except Exception as e: + self.logger.exception(e) + + def _bulk_publish_timed(self, buffer, channel=MQ_CHANNEL, serializer=msgpack.dumps): + total = 0 + try: + for m in buffer: + self.producer.push(metadata=m) + total += len(str(m).encode()) + + except Exception as e: + self.logger.exception(e) + self.logger.error("Some messages couldn't be flushed! Check the messages' contents!") + self.logger.error(f"Message that caused error: {buffer}") + try: + t1 = time() + self.producer.flush() + t2 = time() + self._flush_events.append(["bulk", t1, t2, t2 - t1, total]) + # self.logger.info(f"Flushed {len(buffer)} msgs to MQ!") + except Exception as e: + self.logger.exception(e) + + def liveness_test(self): + """Test Diaspora Liveness.""" + return True + + def unsubscribe(self): + """Unsubscribes from Diaspora topic.""" + raise NotImplementedError() diff --git a/src/flowcept/configs.py b/src/flowcept/configs.py index 204cb7d5..25188ed9 100644 --- a/src/flowcept/configs.py +++ b/src/flowcept/configs.py @@ -26,6 +26,7 @@ "db_buffer": {}, "databases": {"mongodb": {"enabled": False}, "lmdb": {"enabled": False}}, "adapters": {}, + "plugins": {}, "agent": {}, } @@ -298,6 +299,12 @@ def _get_env_list(name: str, default: list[str]) -> list[str]: "MCP_ALLOWED_ORIGINS", AGENT.get("mcp_allowed_origins", ["http://localhost:*", "http://127.0.0.1:*"]) ) +#################### +# PLUGINS # +#################### +# Settings are already resolved to plain containers at load time. +PLUGINS = dict(settings.get("plugins", {}) or {}) + #################### # Enabled ADAPTERS # #################### diff --git a/src/flowcept/flowcept_api/flowcept_controller.py b/src/flowcept/flowcept_api/flowcept_controller.py index 4f243bbc..6ee79fba 100644 --- a/src/flowcept/flowcept_api/flowcept_controller.py +++ b/src/flowcept/flowcept_api/flowcept_controller.py @@ -33,6 +33,7 @@ DUMP_BUFFER_PATH, APPEND_WORKFLOW_ID_TO_PATH, APPEND_ID_TO_PATH, + PLUGINS, ) from flowcept.flowceptor.adapters.base_interceptor import BaseInterceptor @@ -174,6 +175,7 @@ def __init__( self._interceptors = [interceptors] self._interceptor_instances = None + self._plugin_instances: dict = {} self._first_interceptor: BaseInterceptor = None self._should_save_workflow = save_workflow self._current_workflow_obj: WorkflowObject = None @@ -201,6 +203,38 @@ def __init__( if should_delete_buffer_file: Flowcept.delete_buffer_file() + @property + def plugins(self) -> dict: + """Return the dict of running plugin instances keyed by their config name.""" + return self._plugin_instances + + @staticmethod + def _build_plugin(kind: str, cfg: dict): + """Instantiate a plugin by kind string, passing cfg as its config dict. + + Each plugin's internal interceptor automatically reuses Flowcept's + InstrumentationInterceptor singleton (same MQ connection and buffer) when + it is already started — the same pattern used by the Dask client interceptor. + """ + if kind == "academy": + from flowcept.agents.academy.academy_plugin import FlowceptAcademyPlugin + + return FlowceptAcademyPlugin(config=cfg) + elif kind == "langgraph": + from flowcept.agents.langgraph.langgraph_plugin import FlowceptLangGraphPlugin + + return FlowceptLangGraphPlugin(config=cfg) + elif kind == "crewai": + from flowcept.agents.crewai.crewai_plugin import FlowceptCrewAIPlugin + + return FlowceptCrewAIPlugin(config=cfg) + elif kind == "autogen": + from flowcept.agents.autogen.autogen_plugin import FlowceptAutoGenPlugin + + return FlowceptAutoGenPlugin(config=cfg) + else: + raise ValueError(f"Unknown plugin kind: '{kind}'. Supported: academy, langgraph, crewai, autogen.") + def start(self) -> "Flowcept": """Start Flowcept Controller.""" if self.is_started or not self.enabled: @@ -244,6 +278,22 @@ def start(self) -> "Flowcept": else: Flowcept.current_workflow_id = None + + for plugin_name, plugin_cfg in PLUGINS.items(): + if not isinstance(plugin_cfg, dict) or not plugin_cfg.get("enabled", False): + continue + kind = plugin_cfg.get("kind", plugin_name) + # Propagate Flowcept's campaign_id so all auto-started plugins share the same campaign. + merged_cfg = dict(plugin_cfg) + merged_cfg.setdefault("campaign_id", self.campaign_id) + try: + plugin = Flowcept._build_plugin(kind, merged_cfg) + plugin.start() + self._plugin_instances[plugin_name] = plugin + self.logger.debug(f"Plugin '{plugin_name}' ({kind}) started from config.") + except Exception as e: + self.logger.error(f"Failed to start plugin '{plugin_name}': {e}") + Flowcept._current_instance = self Flowcept.is_started = self.is_started = True self.logger.debug("Flowcept started successfully.") @@ -946,6 +996,14 @@ def stop(self): self.logger.warning("Flowcept is already stopped or may never have been started!") return + for plugin_name, plugin in list(self._plugin_instances.items()): + try: + plugin.stop() + self.logger.debug(f"Plugin '{plugin_name}' stopped.") + except Exception as e: + self.logger.error(f"Failed to stop plugin '{plugin_name}': {e}") + self._plugin_instances = {} + if ( self._should_save_workflow and self._first_interceptor is not None diff --git a/src/flowcept/flowceptor/telemetry_capture.py b/src/flowcept/flowceptor/telemetry_capture.py index 3f44061a..36ac4bc2 100644 --- a/src/flowcept/flowceptor/telemetry_capture.py +++ b/src/flowcept/flowceptor/telemetry_capture.py @@ -338,7 +338,7 @@ def _capture_process_info(self): p.executable = psutil_p.exe() p.cmd_line = psutil_p.cmdline() p.num_open_file_descriptors = psutil_p.num_fds() - p.num_connections = len(psutil_p.net_connections()) + p.num_connections = len(psutil_p.connections()) try: p.io_counters = psutil_p.io_counters()._asdict() except Exception: diff --git a/src/flowcept/version.py b/src/flowcept/version.py index bbabfd06..b280f13d 100644 --- a/src/flowcept/version.py +++ b/src/flowcept/version.py @@ -10,4 +10,4 @@ # ❗❗❗ Once again: DO NOT CHANGE THIS FILE ❗❗❗ # ✋⚠️⛔❗❗❗ STOP! DANGER!!ONEONEELEVEN! Did you carefully read the warning above?! :) -__version__ = "1.0.0" +__version__ = "1.0.4" diff --git a/tests/adapters/test_diaspora.py b/tests/adapters/test_diaspora.py new file mode 100644 index 00000000..df8253e4 --- /dev/null +++ b/tests/adapters/test_diaspora.py @@ -0,0 +1,41 @@ +import threading +import unittest + +from flowcept.commons.daos.mq_dao.mq_dao_diaspora import MQDaoDiaspora + + +class TestMQDaoDiaspora(unittest.TestCase): + def setUp(self): + self.dao = MQDaoDiaspora(with_producer=True) + + def test_liveness(self): + self.assertTrue(self.dao.liveness_test()) + + def test_send_and_receive_message(self): + received = [] + + self.dao.subscribe() + + msg = {"task_id": "test-123", "status": "finished"} + + def handler(message): + if message == msg: + received.append(message) + return False # stop after finding our message + return True # keep consuming stale messages + + listener_thread = threading.Thread( + target=self.dao.message_listener, args=(handler,), daemon=True + ) + listener_thread.start() + + self.dao.send_message(msg) + + listener_thread.join(timeout=10) + + self.assertEqual(len(received), 1) + self.assertEqual(received[0], msg) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/agents/__init__.py b/tests/agents/__init__.py new file mode 100644 index 00000000..2b4ed660 --- /dev/null +++ b/tests/agents/__init__.py @@ -0,0 +1 @@ +"""Tests for agentic provenance plugins.""" diff --git a/tests/agents/plugins/__init__.py b/tests/agents/plugins/__init__.py new file mode 100644 index 00000000..2b4ed660 --- /dev/null +++ b/tests/agents/plugins/__init__.py @@ -0,0 +1 @@ +"""Tests for agentic provenance plugins.""" diff --git a/tests/agents/plugins/test_academy_plugin.py b/tests/agents/plugins/test_academy_plugin.py new file mode 100644 index 00000000..ced7212d --- /dev/null +++ b/tests/agents/plugins/test_academy_plugin.py @@ -0,0 +1,323 @@ +"""Unit tests for the Academy provenance plugin. + +These tests never touch MongoDB, Redis, or the network: the FlowCept +``BaseInterceptor`` is replaced by an in-memory fake that appends every +emitted workflow/task record to a plain list, while the plugin's own logic +(enrichment, patching of ``academy.runtime.Runtime``, ContextVars, LLM hook) +runs for real against a local Academy exchange. +""" + +from __future__ import annotations + +import asyncio +from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor + +import pytest + +pytest.importorskip("academy") + +from academy.agent import Agent, action # noqa: E402 +from academy.exchange import LocalExchangeFactory # noqa: E402 +from academy.manager import Manager # noqa: E402 + +import flowcept.flowceptor.adapters.base_interceptor as bi_mod # noqa: E402 +from flowcept.agents.academy import academy_plugin as ap # noqa: E402 +from flowcept.agents.academy.academy_plugin import FlowceptAcademyPlugin # noqa: E402 + + +# -- capture fixture ---------------------------------------------------------- + + +@pytest.fixture +def capture(monkeypatch): + """Replace BaseInterceptor with an in-memory fake and return the record list.""" + state = {"records": [], "stop_calls": 0} + + class _FakeBaseInterceptor: + """In-memory stand-in for FlowCept's BaseInterceptor (no MQ, no DB).""" + + def __init__(self, plugin_key=None, kind=None): + self.kind = kind + self.telemetry_capture = None + self.started = False + + def start(self, bundle_exec_id, check_safe_stops=True): + """Mark the interceptor as started.""" + self.started = True + return self + + def stop(self, check_safe_stops=True): + """Count flushes instead of talking to an MQ.""" + state["stop_calls"] += 1 + self.started = False + + def intercept(self, obj): + """Append a task record to the shared list.""" + state["records"].append(obj) + + def send_workflow_message(self, wf): + """Append a workflow record to the shared list.""" + state["records"].append(wf.to_dict()) + + monkeypatch.setattr(bi_mod, "BaseInterceptor", _FakeBaseInterceptor) + yield state + ap._ACTIVE_INTERCEPTOR = None + + +@pytest.fixture +def plugin(capture): + """Return a started plugin wired to the in-memory capture fixture.""" + p = FlowceptAcademyPlugin(config={"enabled": True, "workflow_name": "academy-test", "performance_tracking": False}) + p.start() + yield p + p.stop() + + +def _tasks(records, subtype=None): + """Return captured task records, optionally filtered by subtype.""" + return [r for r in records if r.get("type") == "task" and (subtype is None or r.get("subtype") == subtype)] + + +def _workflows(records): + """Return captured workflow records.""" + return [r for r in records if r.get("type") == "workflow"] + + +# -- test agent --------------------------------------------------------------- + + +class EchoAgent(Agent): + """Minimal Academy agent exercising success, failure, and ContextVars.""" + + @action + async def double(self, n: int) -> int: + """Return twice the input.""" + return 2 * n + + @action + async def boom(self) -> None: + """Raise ValueError to exercise the error path.""" + raise ValueError("nope") + + @action + async def whoami(self) -> tuple: + """Return the cross-framework linking ContextVar values.""" + return (ap._current_action_task_id.get(), ap._current_academy_agent_id.get()) + + @action + async def call_llm(self) -> None: + """Record a synthetic LLM call from inside an action.""" + ap.record_llm_call( + { + "type": "chat_completion", + "model": "fake-model", + "text": '{"score": 5}', + "finish_reason": "stop", + "usage": {"prompt_tokens": 3, "completion_tokens": 4, "total_tokens": 7}, + "context": {"call_type": "scoring"}, + } + ) + + +def _run_agent(script): + """Launch one EchoAgent on a local exchange and run *script(handle)*.""" + + async def main(): + factory = LocalExchangeFactory() + executor = ThreadPoolExecutor(max_workers=2) + async with await Manager.from_exchange_factory(factory=factory, executors=executor) as manager: + handle = await manager.launch(EchoAgent) + return await script(handle) + + return asyncio.run(main()) + + +# -- lifecycle ---------------------------------------------------------------- + + +def test_plugin_is_disabled_by_default(capture): + """Without enabled=True in config, start() must be a no-op.""" + p = FlowceptAcademyPlugin(config={"workflow_name": "x"}) + p.start() + assert p._started is False + assert capture["records"] == [] + + +def test_start_emits_top_level_workflow(plugin, capture): + """Starting the plugin emits one WorkflowObject with the configured name.""" + wfs = _workflows(capture["records"]) + assert len(wfs) == 1 + assert wfs[0]["name"] == "academy-test" + assert wfs[0]["workflow_id"] == plugin._interceptor._workflow_id + assert wfs[0]["campaign_id"] == plugin._interceptor._campaign_id + + +def test_start_uses_configured_campaign_id(capture): + """An explicit campaign_id in config is propagated to every record.""" + p = FlowceptAcademyPlugin(config={"enabled": True, "campaign_id": "camp-42", "performance_tracking": False}) + p.start() + try: + p._interceptor.intercept_task({"activity_id": "a"}) + assert _workflows(capture["records"])[0]["campaign_id"] == "camp-42" + assert _tasks(capture["records"])[0]["campaign_id"] == "camp-42" + finally: + p.stop() + + +def test_start_registers_llm_hook_and_stop_unregisters(capture): + """The LLM hook register/unregister callables receive the plugin's hook.""" + registered, unregistered = [], [] + p = FlowceptAcademyPlugin( + config={"enabled": True, "performance_tracking": False}, + llm_hook_register=registered.append, + llm_hook_unregister=unregistered.append, + ) + p.start() + p.stop() + assert registered == [ap._on_llm_call] + assert unregistered == [ap._on_llm_call] + + +def test_stop_flushes_interceptor_and_clears_active_state(plugin, capture): + """stop() flushes the interceptor and deactivates the module-level state.""" + plugin.stop() + assert capture["stop_calls"] == 1 + assert plugin._started is False + assert ap._ACTIVE_INTERCEPTOR is None + + +def test_stop_restores_process_pool_executor(capture): + """start() patches ProcessPoolExecutor.__init__ and stop() restores it.""" + original = ProcessPoolExecutor.__init__ + p = FlowceptAcademyPlugin(config={"enabled": True, "performance_tracking": False}) + p.start() + assert ProcessPoolExecutor.__init__ is not original + p.stop() + assert ProcessPoolExecutor.__init__ is original + + +# -- interceptor enrichment --------------------------------------------------- + + +def test_intercept_task_fills_standard_fields(plugin, capture): + """intercept_task adds type, ids, and normalizes the status enum value.""" + plugin._interceptor.intercept_task({"activity_id": "my_act", "status": "FINISHED"}) + task = _tasks(capture["records"])[0] + assert task["type"] == "task" + assert task["task_id"] + assert task["workflow_id"] == plugin._interceptor._workflow_id + assert task["campaign_id"] == plugin._interceptor._campaign_id + assert task["status"] == "FINISHED" + # enrich_task_dict adds host identity fields + assert "hostname" in task + + +def test_intercept_task_normalizes_unknown_status(plugin, capture): + """An unrecognized status string falls back to FINISHED.""" + plugin._interceptor.intercept_task({"activity_id": "a", "status": "banana"}) + assert _tasks(capture["records"])[0]["status"] == "FINISHED" + + +# -- real Academy runs -------------------------------------------------------- + + +def test_action_run_emits_finished_task(plugin, capture): + """A successful @action produces an academy_action task with used/generated.""" + + async def script(handle): + return await handle.double(21) + + assert _run_agent(script) == 42 + actions = [t for t in _tasks(capture["records"], "academy_action") if t["activity_id"] == "double"] + assert len(actions) == 1 + task = actions[0] + assert task["status"] == "FINISHED" + assert task["used"]["args"] == [21] + assert task["generated"] == 42 + assert task["ended_at"] >= task["started_at"] + assert task["custom_metadata"]["agent_type"] == "EchoAgent" + assert task["custom_metadata"]["cross_agent_call"] is False + + +def test_action_failure_emits_error_task(plugin, capture): + """A raising @action produces a task with status ERROR and the stderr text.""" + + async def script(handle): + with pytest.raises(Exception): + await handle.boom() + + _run_agent(script) + task = [t for t in _tasks(capture["records"], "academy_action") if t["activity_id"] == "boom"][0] + assert task["status"] == "ERROR" + assert "nope" in task["stderr"] + assert task["generated"] is None + + +def test_agent_lifecycle_records_and_sub_workflow(plugin, capture): + """Agent startup/shutdown emit lifecycle tasks plus a linked sub-workflow.""" + + async def script(handle): + return await handle.double(1) + + _run_agent(script) + lifecycle = _tasks(capture["records"], "academy_lifecycle") + events = [t["activity_id"] for t in lifecycle] + assert "agent_startup" in events + assert "agent_shutdown" in events + + top_wf_id = plugin._interceptor._workflow_id + sub_wfs = [w for w in _workflows(capture["records"]) if w.get("parent_workflow_id")] + assert len(sub_wfs) == 1 + assert sub_wfs[0]["parent_workflow_id"] == top_wf_id + assert sub_wfs[0]["custom_metadata"]["agent_type"] == "EchoAgent" + + +def test_contextvars_expose_action_task_id_and_agent_id(plugin, capture): + """Inside an @action, the cross-framework linking ContextVars are set.""" + + async def script(handle): + return await handle.whoami() + + action_task_id, academy_agent_id = _run_agent(script) + assert action_task_id is not None + assert academy_agent_id is not None and academy_agent_id.startswith("AgentId") + + task = [t for t in _tasks(capture["records"], "academy_action") if t["activity_id"] == "whoami"][0] + assert task["task_id"] == action_task_id + assert task["agent_id"] == academy_agent_id + # Outside any action, the ContextVars are unset in this context. + assert ap._current_action_task_id.get() is None + + +def test_llm_call_inside_action_links_parent_task(plugin, capture): + """record_llm_call() inside an @action becomes a child llm_call task.""" + + async def script(handle): + return await handle.call_llm() + + _run_agent(script) + action_task = [t for t in _tasks(capture["records"], "academy_action") if t["activity_id"] == "call_llm"][0] + llm = _tasks(capture["records"], "llm_call")[0] + assert llm["parent_task_id"] == action_task["task_id"] + assert llm["activity_id"] == "scoring" + assert llm["agent_id"] == action_task["agent_id"] + assert llm["used"]["model"] == "fake-model" + assert llm["generated"]["total_tokens"] == 7 + # JSON embedded in the response text is parsed and hoisted for queries. + assert llm["generated"]["parsed_response"] == {"score": 5} + assert llm["generated"]["score"] == 5 + + +def test_record_llm_call_noop_when_plugin_not_started(capture): + """record_llm_call() must not emit or raise when no plugin is active.""" + assert ap._ACTIVE_INTERCEPTOR is None + ap.record_llm_call({"type": "chat_completion", "model": "m", "text": "x"}) + assert capture["records"] == [] + + +def test_llm_call_error_payload_marks_task_failed(plugin, capture): + """An LLM payload containing 'error' produces an ERROR llm_call task.""" + ap.record_llm_call({"type": "chat_completion", "model": "m", "error": "rate limited"}) + llm = _tasks(capture["records"], "llm_call")[0] + assert llm["status"] == "ERROR" + assert llm["generated"]["error"] == "rate limited" diff --git a/tests/agents/plugins/test_autogen_plugin.py b/tests/agents/plugins/test_autogen_plugin.py new file mode 100644 index 00000000..41032205 --- /dev/null +++ b/tests/agents/plugins/test_autogen_plugin.py @@ -0,0 +1,338 @@ +"""Unit tests for the AutoGen provenance plugin. + +These tests never touch MongoDB, Redis, an LLM API, or the network: FlowCept's +``BaseInterceptor`` is replaced by an in-memory fake that appends every +emitted workflow/task record to a plain list, and teams/model clients are +lightweight fakes shaped like the real AutoGen payloads (real +``TaskResult``/``TextMessage`` objects are used where the plugin type-checks). +""" + +from __future__ import annotations + +import asyncio + +import pytest + +pytest.importorskip("autogen_agentchat") +pytest.importorskip("autogen_core") + +from autogen_agentchat.agents import AssistantAgent # noqa: E402 +from autogen_agentchat.base import TaskResult # noqa: E402 +from autogen_agentchat.messages import TextMessage # noqa: E402 + +import flowcept.flowceptor.adapters.base_interceptor as bi_mod # noqa: E402 +from flowcept.agents.autogen import autogen_plugin as agp # noqa: E402 +from flowcept.agents.autogen.autogen_plugin import ( # noqa: E402 + FlowceptAutoGenPlugin, + FlowceptModelClient, +) + + +# -- capture fixture ---------------------------------------------------------- + + +@pytest.fixture +def capture(monkeypatch): + """Replace BaseInterceptor with an in-memory fake and return the record list.""" + state = {"records": [], "stop_calls": 0} + + class _FakeBaseInterceptor: + """In-memory stand-in for FlowCept's BaseInterceptor (no MQ, no DB).""" + + def __init__(self, plugin_key=None, kind=None): + self.kind = kind + self.telemetry_capture = None + + def start(self, bundle_exec_id, check_safe_stops=True): + """Pretend to start.""" + return self + + def stop(self, check_safe_stops=True): + """Count flushes instead of talking to an MQ.""" + state["stop_calls"] += 1 + + def intercept(self, obj): + """Append a task record to the shared list.""" + state["records"].append(obj) + + def send_workflow_message(self, wf): + """Append a workflow record to the shared list.""" + state["records"].append(wf.to_dict()) + + monkeypatch.setattr(bi_mod, "BaseInterceptor", _FakeBaseInterceptor) + yield state + agp._ACTIVE_INTERCEPTOR = None + agp._PROV_STATS = None + + +@pytest.fixture +def plugin(capture): + """Return a started plugin wired to the in-memory capture fixture.""" + p = FlowceptAutoGenPlugin(config={"workflow_name": "autogen-test", "performance_tracking": False}) + p.start() + yield p + p.stop() + + +def _tasks(records, subtype=None): + """Return captured task records, optionally filtered by subtype.""" + return [r for r in records if "subtype" in r and (subtype is None or r.get("subtype") == subtype)] + + +def _workflows(records): + """Return captured workflow records.""" + return [r for r in records if r.get("type") == "workflow"] + + +# -- fakes -------------------------------------------------------------------- + + +class FakeTeam: + """Minimal team exposing run_stream(), shaped like an AutoGen group chat.""" + + name = "fake_team" + _participants: list = [] + + def __init__(self, fail_after=None, llm_between_messages=False): + self._fail_after = fail_after + self._llm_between_messages = llm_between_messages + + async def run_stream(self, task, cancellation_token=None): + """Yield two chat messages then a TaskResult, optionally failing midway.""" + yield TextMessage(source="user", content=task) + if self._fail_after == 1: + raise RuntimeError("team exploded") + if self._llm_between_messages: + agp.record_llm_call( + { + "type": "chat_completion", + "model": "fake-model", + "text": "answer", + "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}, + } + ) + yield TextMessage(source="agent1", content="answer") + yield TaskResult(messages=[], stop_reason="done") + + +class FakeChatCompletionClient: + """Duck-typed AutoGen ChatCompletionClient returning a canned response.""" + + model_info = { + "vision": False, + "function_calling": False, + "json_output": False, + "structured_output": False, + "family": "fake", + } + model = "fake-model" + + async def create(self, messages, **kwargs): + """Return a canned CreateResult-shaped object.""" + + class _Usage: + prompt_tokens = 3 + completion_tokens = 5 + + class _Result: + content = "hi there" + finish_reason = "stop" + usage = _Usage() + + return _Result() + + async def close(self): + """No-op close.""" + + +# -- lifecycle ---------------------------------------------------------------- + + +def test_start_emits_top_level_workflow(plugin, capture): + """Starting the plugin emits one WorkflowObject with the configured name.""" + wfs = _workflows(capture["records"]) + assert len(wfs) == 1 + assert wfs[0]["name"] == "autogen-test" + assert wfs[0]["workflow_id"] == plugin._interceptor._workflow_id + assert wfs[0]["campaign_id"] == plugin._interceptor._campaign_id + + +def test_disabled_plugin_does_not_start(capture): + """With enabled=False in config, start() must be a no-op.""" + p = FlowceptAutoGenPlugin(config={"enabled": False}) + p.start() + assert p._started is False + assert capture["records"] == [] + + +def test_context_manager_starts_and_stops(capture): + """The plugin works as a context manager, flushing on exit.""" + with FlowceptAutoGenPlugin(config={"performance_tracking": False}) as p: + assert p._started is True + assert agp._ACTIVE_INTERCEPTOR is p._interceptor + assert p._started is False + assert capture["stop_calls"] == 1 + assert agp._ACTIVE_INTERCEPTOR is None + + +def test_stop_restores_assistant_agent_init(capture): + """start() patches AssistantAgent.__init__ and stop() restores it.""" + original = AssistantAgent.__init__ + p = FlowceptAutoGenPlugin(config={"performance_tracking": False}) + p.start() + assert AssistantAgent.__init__ is not original + p.stop() + assert AssistantAgent.__init__ is original + + +def test_intercept_task_fills_standard_fields(plugin, capture): + """intercept_task adds ids and normalizes the status enum value.""" + plugin._interceptor.intercept_task({"activity_id": "a", "status": "ERROR"}) + task = capture["records"][-1] + assert task["activity_id"] == "a" + assert task["task_id"] + assert task["workflow_id"] == plugin._interceptor._workflow_id + assert task["campaign_id"] == plugin._interceptor._campaign_id + assert task["status"] == "ERROR" + assert "hostname" in task + + +# -- team runs ---------------------------------------------------------------- + + +def test_run_team_emits_run_and_message_tasks(plugin, capture): + """A team run yields one autogen_run task plus one task per message.""" + result = asyncio.run(plugin.run_team(FakeTeam(), "hello")) + assert isinstance(result, TaskResult) + assert result.stop_reason == "done" + + run = _tasks(capture["records"], "autogen_run")[0] + msgs = _tasks(capture["records"], "autogen_message") + assert run["activity_id"] == "fake_team" + assert run["status"] == "FINISHED" + assert run["used"] == {"task": "hello"} + assert run["generated"]["stop_reason"] == "done" + assert run["generated"]["message_count"] == 2 + assert [m["activity_id"] for m in msgs] == ["user", "agent1"] + assert msgs[0]["generated"]["content"] == "hello" + assert msgs[1]["generated"]["content"] == "answer" + assert msgs[1]["generated"]["message_type"] == "TextMessage" + + +def test_message_tasks_share_group_id_and_parent(plugin, capture): + """All messages of one run share group_id and are children of the run task.""" + asyncio.run(plugin.run_team(FakeTeam(), "hello")) + run = _tasks(capture["records"], "autogen_run")[0] + msgs = _tasks(capture["records"], "autogen_message") + assert {m["group_id"] for m in msgs} == {run["group_id"]} + assert {m["parent_task_id"] for m in msgs} == {run["task_id"]} + # The run also emits a sub-workflow linked to the top-level workflow. + sub_wfs = [w for w in _workflows(capture["records"]) if w.get("parent_workflow_id")] + assert sub_wfs[0]["parent_workflow_id"] == plugin._interceptor._workflow_id + assert sub_wfs[0]["custom_metadata"]["group_id"] == run["group_id"] + + +def test_run_team_records_cross_framework_source_agent_id(plugin, capture): + """A source_agent_id from another framework lands in custom_metadata.""" + asyncio.run(plugin.run_team(FakeTeam(), "hello", source_agent_id="AgentId")) + run = _tasks(capture["records"], "autogen_run")[0] + assert run["custom_metadata"]["source_agent_id"] == "AgentId" + + +def test_run_team_failure_emits_error_run_record(plugin, capture): + """A failure mid-stream re-raises and records the run task as ERROR.""" + with pytest.raises(RuntimeError, match="team exploded"): + asyncio.run(plugin.run_team(FakeTeam(fail_after=1), "hello")) + run = _tasks(capture["records"], "autogen_run")[0] + assert run["status"] == "ERROR" + assert "team exploded" in run["stderr"] + # The message seen before the failure was still captured. + assert len(_tasks(capture["records"], "autogen_message")) == 1 + + +def test_run_team_without_start_runs_plain(capture): + """An unstarted plugin still runs the team but emits no provenance.""" + p = FlowceptAutoGenPlugin(config={"enabled": False}) + result = asyncio.run(p.run_team(FakeTeam(), "hello")) + assert isinstance(result, TaskResult) + assert capture["records"] == [] + + +def test_module_level_run_team_uses_active_interceptor(plugin, capture): + """agp.run_team() picks up the interceptor started by the plugin.""" + result = asyncio.run(agp.run_team(FakeTeam(), "hello", team_name="named_run")) + assert isinstance(result, TaskResult) + run = _tasks(capture["records"], "autogen_run")[0] + assert run["activity_id"] == "named_run" + + +# -- LLM call capture --------------------------------------------------------- + + +def test_record_llm_call_emits_llm_task(plugin, capture): + """record_llm_call() builds an llm_call task with used/generated payloads.""" + agp.record_llm_call( + { + "type": "chat_completion", + "model": "fake-model", + "user_prompt": "q", + "temperature": 0.1, + "text": "a", + "finish_reason": "stop", + "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}, + "context": {"agent_name": "agent1"}, + } + ) + llm = _tasks(capture["records"], "llm_call")[0] + assert llm["activity_id"] == "fake-model" + assert llm["status"] == "FINISHED" + assert llm["used"]["user_prompt"] == "q" + assert llm["used"]["temperature"] == 0.1 + assert llm["generated"]["text"] == "a" + assert llm["generated"]["usage"]["total_tokens"] == 3 + assert llm["custom_metadata"]["agent_id"] == "agent1" + + +def test_record_llm_call_error_marks_task_failed(plugin, capture): + """A payload containing 'error' produces an ERROR llm_call task.""" + agp.record_llm_call({"type": "chat_completion", "model": "m", "error": "rate limited"}) + llm = _tasks(capture["records"], "llm_call")[0] + assert llm["status"] == "ERROR" + assert llm["generated"]["error"] == "rate limited" + + +def test_record_llm_call_noop_when_plugin_not_started(capture): + """record_llm_call() must not emit or raise when no plugin is active.""" + assert agp._ACTIVE_INTERCEPTOR is None + agp.record_llm_call({"type": "chat_completion", "model": "m", "text": "x"}) + assert capture["records"] == [] + + +def test_llm_call_inside_stream_links_to_message_task(plugin, capture): + """An LLM call made mid-stream is a child of the current message task.""" + asyncio.run(plugin.run_team(FakeTeam(llm_between_messages=True), "hello")) + first_msg = _tasks(capture["records"], "autogen_message")[0] + llm = _tasks(capture["records"], "llm_call")[0] + assert llm["parent_task_id"] == first_msg["task_id"] + assert llm["custom_metadata"]["agent_id"] == first_msg["activity_id"] + + +def test_flowcept_model_client_records_usage(plugin, capture): + """FlowceptModelClient.create() records model, agent, and token usage.""" + wrapped = FlowceptModelClient(FakeChatCompletionClient(), agent_name="worker") + result = asyncio.run(wrapped.create([TextMessage(source="u", content="q")], temperature=0.2)) + assert result.content == "hi there" + + llm = _tasks(capture["records"], "llm_call")[0] + assert llm["activity_id"] == "fake-model" + assert llm["used"]["temperature"] == 0.2 + assert llm["generated"]["text"] == "hi there" + assert llm["generated"]["usage"] == {"prompt_tokens": 3, "completion_tokens": 5, "total_tokens": 8} + assert llm["custom_metadata"]["agent_id"] == "worker" + assert llm["custom_metadata"]["framework"] == "autogen" + + +def test_assistant_agent_model_client_auto_wrapped(plugin): + """While the plugin runs, new AssistantAgents get a wrapped model client.""" + agent = AssistantAgent(name="a1", model_client=FakeChatCompletionClient()) + assert isinstance(agent._model_client, FlowceptModelClient) diff --git a/tests/agents/plugins/test_crewai_plugin.py b/tests/agents/plugins/test_crewai_plugin.py new file mode 100644 index 00000000..6e39c420 --- /dev/null +++ b/tests/agents/plugins/test_crewai_plugin.py @@ -0,0 +1,387 @@ +"""Unit tests for the FlowCept CrewAI provenance plugin. + +Provenance emission is captured in memory by replacing ``BaseInterceptor`` +with a fake that records every workflow and task message, so no MQ, MongoDB, +or network access is needed. Because running a real Crew requires an LLM, +these tests drive the plugin's event-bus listener and LLM/tool hook surfaces +directly with synthetic events, mirroring the direct-handler-call style used +elsewhere in the test suite. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +import flowcept.flowceptor.adapters.base_interceptor as base_interceptor_module +from flowcept.agents.crewai.crewai_plugin import FlowceptCrewAIPlugin + +pytest.importorskip("crewai") + + +class _CapturingInterceptor: + """In-memory stand-in for BaseInterceptor that records all emissions.""" + + instances: list = [] + + def __init__(self, plugin_key=None, kind=None): + """Record construction and initialize empty capture buffers.""" + self.kind = kind + self.telemetry_capture = None + self.started = False + self.stopped = False + self.task_messages: list[dict] = [] + self.workflow_messages: list = [] + type(self).instances.append(self) + + def start(self, bundle_exec_id, check_safe_stops=True): + """Mark the interceptor as started.""" + self.started = True + return self + + def stop(self, check_safe_stops=True): + """Mark the interceptor as stopped.""" + self.stopped = True + + def send_workflow_message(self, workflow_obj): + """Capture a WorkflowObject emission.""" + self.workflow_messages.append(workflow_obj) + + def intercept(self, task_dict): + """Capture a task message emission.""" + self.task_messages.append(task_dict) + + +@pytest.fixture() +def captured(monkeypatch): + """Patch BaseInterceptor with the capturing fake and return its instance list.""" + _CapturingInterceptor.instances = [] + monkeypatch.setattr(base_interceptor_module, "BaseInterceptor", _CapturingInterceptor) + return _CapturingInterceptor.instances + + +@pytest.fixture() +def started_plugin(captured): + """Yield a started plugin plus the fake interceptor backing it.""" + plugin = FlowceptCrewAIPlugin(config={"workflow_name": "crew-test-wf", "performance_tracking": False}) + plugin.start() + assert captured, "plugin.start() did not build an interceptor" + yield plugin, captured[-1] + plugin.stop() + + +def _run_crew_lifecycle(listener, crew_name="my-crew", inputs=None, fail=False): + """Drive a full synthetic crew kickoff through the listener callbacks.""" + started = SimpleNamespace(crew_name=crew_name, inputs=inputs or {}, event_id="crew-evt-1") + listener.on_crew_kickoff_started(None, started) + if fail: + failed = SimpleNamespace(event_id="crew-evt-2", started_event_id="crew-evt-1", error="kickoff exploded") + listener.on_crew_kickoff_failed(None, failed) + else: + completed = SimpleNamespace( + event_id="crew-evt-2", started_event_id="crew-evt-1", output="crew output", total_tokens=42 + ) + listener.on_crew_kickoff_completed(None, completed) + + +# -- lifecycle ---------------------------------------------------------------- + + +def test_start_emits_top_level_workflow_message(started_plugin): + """start() sends one WorkflowObject carrying the configured workflow name.""" + _, fake = started_plugin + assert len(fake.workflow_messages) == 1 + wf = fake.workflow_messages[0] + assert wf.name == "crew-test-wf" + assert wf.workflow_id is not None + assert wf.campaign_id is not None + + +def test_start_registers_llm_and_tool_hooks(started_plugin): + """start() registers the plugin's global CrewAI LLM and tool hooks.""" + from crewai.hooks.llm_hooks import get_after_llm_call_hooks, get_before_llm_call_hooks + from crewai.hooks.tool_hooks import get_after_tool_call_hooks, get_before_tool_call_hooks + + plugin, _ = started_plugin + assert plugin._hooks_obj.before_llm_call in get_before_llm_call_hooks() + assert plugin._hooks_obj.after_llm_call in get_after_llm_call_hooks() + assert plugin._hooks_obj.before_tool_call in get_before_tool_call_hooks() + assert plugin._hooks_obj.after_tool_call in get_after_tool_call_hooks() + + +def test_stop_unregisters_hooks_and_stops_interceptor(started_plugin): + """stop() removes the global hooks and stops the wrapped interceptor.""" + from crewai.hooks.llm_hooks import get_before_llm_call_hooks + from crewai.hooks.tool_hooks import get_before_tool_call_hooks + + plugin, fake = started_plugin + plugin.stop() + assert fake.stopped is True + assert plugin._hooks_obj.before_llm_call not in get_before_llm_call_hooks() + assert plugin._hooks_obj.before_tool_call not in get_before_tool_call_hooks() + plugin.stop() # second stop is a safe no-op + + +def test_disabled_plugin_emits_nothing(captured): + """enabled=False disables capture entirely and never builds an interceptor.""" + plugin = FlowceptCrewAIPlugin(config={"enabled": False}) + plugin.start() + assert captured == [] + plugin.stop() + + +def test_context_manager_starts_and_stops(captured): + """The plugin works as a context manager, starting on enter and stopping on exit.""" + with FlowceptCrewAIPlugin(config={"workflow_name": "ctx-crew", "performance_tracking": False}) as plugin: + assert plugin._started is True + fake = captured[-1] + assert plugin._started is False + assert fake.stopped is True + + +def test_start_respects_custom_campaign_id(captured): + """A campaign_id passed in config is used verbatim on the workflow message.""" + plugin = FlowceptCrewAIPlugin( + config={"workflow_name": "wf", "campaign_id": "camp-7", "performance_tracking": False} + ) + plugin.start() + try: + assert captured[-1].workflow_messages[0].campaign_id == "camp-7" + finally: + plugin.stop() + + +# -- crew kickoff ------------------------------------------------------------- + + +def test_crew_kickoff_emits_crew_task_and_sub_workflow(started_plugin): + """A kickoff start/complete pair yields one crewai_crew task and a sub-workflow.""" + plugin, fake = started_plugin + _run_crew_lifecycle(plugin._listener_obj, inputs={"topic": "prov"}) + + crew_task = next(t for t in fake.task_messages if t["subtype"] == "crewai_crew") + assert crew_task["activity_id"] == "my-crew" + assert crew_task["status"] == "FINISHED" + assert crew_task["used"]["inputs"] == {"topic": "prov"} + assert crew_task["generated"]["output"] == "crew output" + assert crew_task["generated"]["total_tokens"] == 42 + assert crew_task["workflow_id"] == fake.workflow_messages[0].workflow_id + assert crew_task["campaign_id"] == fake.workflow_messages[0].campaign_id + + top_wf, sub_wf = fake.workflow_messages[0], fake.workflow_messages[1] + assert sub_wf.name == "my-crew" + assert sub_wf.parent_workflow_id == top_wf.workflow_id + assert sub_wf.custom_metadata["group_id"] == crew_task["group_id"] + + +def test_crew_kickoff_failure_recorded_as_error(started_plugin): + """A kickoff failure event yields an ERROR crewai_crew task with stderr.""" + plugin, fake = started_plugin + _run_crew_lifecycle(plugin._listener_obj, fail=True) + + crew_task = next(t for t in fake.task_messages if t["subtype"] == "crewai_crew") + assert crew_task["status"] == "ERROR" + assert "kickoff exploded" in crew_task["stderr"] + + +# -- task and agent lifecycle --------------------------------------------------- + + +def test_task_lifecycle_emits_crewai_task(started_plugin): + """A task start/complete pair yields one crewai_task with metadata and output.""" + plugin, fake = started_plugin + listener = plugin._listener_obj + listener.on_crew_kickoff_started(None, SimpleNamespace(crew_name="c", inputs={}, event_id="ck-1")) + listener.on_task_started( + None, + SimpleNamespace(event_id="tk-1", task_name="research", agent_role="researcher", context={"c": 1}, task=None), + ) + listener.on_task_completed(None, SimpleNamespace(event_id="tk-2", started_event_id="tk-1", output="findings")) + + task = next(t for t in fake.task_messages if t["subtype"] == "crewai_task") + assert task["activity_id"] == "research" + assert task["status"] == "FINISHED" + assert task["used"]["context"] == {"c": 1} + assert task["generated"]["output"] == "findings" + assert task["custom_metadata"]["task_name"] == "research" + assert task["custom_metadata"]["agent_role"] == "researcher" + crew_group = next(iter(listener._crew_group.values())) + assert task["group_id"] == crew_group + + +def test_task_failure_recorded_as_error(started_plugin): + """A task failure event yields an ERROR crewai_task with stderr.""" + plugin, fake = started_plugin + listener = plugin._listener_obj + listener.on_task_started( + None, SimpleNamespace(event_id="tk-1", task_name="research", agent_role="r", context={}, task=None) + ) + listener.on_task_failed(None, SimpleNamespace(event_id="tk-2", started_event_id="tk-1", error="task blew up")) + + task = next(t for t in fake.task_messages if t["subtype"] == "crewai_task") + assert task["status"] == "ERROR" + assert "task blew up" in task["stderr"] + + +def test_agent_execution_links_to_enclosing_task(started_plugin): + """A crewai_agent task carries parent_task_id of the enclosing crewai_task.""" + plugin, fake = started_plugin + listener = plugin._listener_obj + listener.on_task_started( + None, SimpleNamespace(event_id="tk-1", task_name="research", agent_role="researcher", context={}, task=None) + ) + listener.on_agent_execution_started( + None, + SimpleNamespace( + event_id="ag-1", + agent=SimpleNamespace(role="researcher"), + agent_role="researcher", + task_prompt="find sources", + tools=[SimpleNamespace(name="search")], + started_event_id="tk-1", + task_id=None, + ), + ) + listener.on_agent_execution_completed( + None, SimpleNamespace(event_id="ag-2", started_event_id="ag-1", output="done") + ) + listener.on_task_completed(None, SimpleNamespace(event_id="tk-2", started_event_id="tk-1", output="out")) + + agent_task = next(t for t in fake.task_messages if t["subtype"] == "crewai_agent") + crew_task = next(t for t in fake.task_messages if t["subtype"] == "crewai_task") + assert agent_task["activity_id"] == "researcher" + assert agent_task["used"]["task_prompt"] == "find sources" + assert agent_task["used"]["tools"] == ["search"] + assert agent_task["generated"]["output"] == "done" + assert agent_task["parent_task_id"] == crew_task["task_id"] + + +def test_agent_execution_error_recorded(started_plugin): + """An agent execution error yields an ERROR crewai_agent task.""" + plugin, fake = started_plugin + listener = plugin._listener_obj + listener.on_agent_execution_started( + None, + SimpleNamespace( + event_id="ag-1", + agent=None, + agent_role="researcher", + task_prompt="x", + tools=[], + started_event_id=None, + task_id=None, + ), + ) + listener.on_agent_execution_error( + None, SimpleNamespace(event_id="ag-2", started_event_id="ag-1", error="agent crashed") + ) + + agent_task = next(t for t in fake.task_messages if t["subtype"] == "crewai_agent") + assert agent_task["status"] == "ERROR" + assert "agent crashed" in agent_task["stderr"] + + +# -- LLM and tool hooks ---------------------------------------------------------- + + +def test_llm_hooks_emit_llm_call_with_agent_context(started_plugin): + """The before/after LLM hooks emit one llm_call task linked to the agent task.""" + plugin, fake = started_plugin + listener, hooks = plugin._listener_obj, plugin._hooks_obj + listener.on_crew_kickoff_started(None, SimpleNamespace(crew_name="c", inputs={}, event_id="ck-1")) + listener.on_agent_execution_started( + None, + SimpleNamespace( + event_id="ag-1", + agent=SimpleNamespace(role="writer"), + agent_role="writer", + task_prompt="write", + tools=[], + started_event_id=None, + task_id=None, + ), + ) + executor = object() + hooks.before_llm_call( + SimpleNamespace( + agent=SimpleNamespace(role="writer", goal="write well", backstory="b"), + task=SimpleNamespace(description="write a poem", expected_output="poem"), + llm=SimpleNamespace(model="fake-model"), + messages=[{"role": "user", "content": "write a poem"}], + iterations=2, + executor=executor, + ) + ) + hooks.after_llm_call(SimpleNamespace(executor=executor, response="roses are red")) + + llm_task = next(t for t in fake.task_messages if t["subtype"] == "llm_call") + agent_fc_id = next(iter(listener._agent_fc_id.values())) + assert llm_task["activity_id"] == "fake-model" + assert llm_task["status"] == "FINISHED" + assert llm_task["used"]["messages"] == [{"role": "user", "content": "write a poem"}] + assert llm_task["used"]["agent_role"] == "writer" + assert llm_task["used"]["task_description"] == "write a poem" + assert llm_task["used"]["iterations"] == 2 + assert llm_task["generated"]["response"] == "roses are red" + assert llm_task["parent_task_id"] == agent_fc_id + assert llm_task["group_id"] == next(iter(listener._crew_group.values())) + assert llm_task["custom_metadata"]["source"] == "llm_hook" + + +def test_tool_hooks_emit_tool_call_with_typed_input(started_plugin): + """The before/after tool hooks emit one tool_call task with the typed input.""" + plugin, fake = started_plugin + hooks = plugin._hooks_obj + executor = object() + hooks.before_tool_call( + SimpleNamespace( + tool_name="web_search", + agent=SimpleNamespace(role="researcher"), + task=SimpleNamespace(description="find sources"), + tool_input={"query": "flowcept"}, + executor=executor, + ) + ) + hooks.after_tool_call(SimpleNamespace(executor=executor, tool_result="found 3 results")) + + tool_task = next(t for t in fake.task_messages if t["subtype"] == "tool_call") + assert tool_task["activity_id"] == "web_search" + assert tool_task["status"] == "FINISHED" + assert tool_task["used"]["input"] == {"query": "flowcept"} + assert tool_task["used"]["agent_role"] == "researcher" + assert tool_task["generated"]["output"] == "found 3 results" + assert tool_task["custom_metadata"]["tool_name"] == "web_search" + + +def test_all_tasks_in_one_kickoff_share_group_id(started_plugin): + """Crew, task, and agent records of one kickoff carry the same group_id.""" + plugin, fake = started_plugin + listener = plugin._listener_obj + listener.on_crew_kickoff_started(None, SimpleNamespace(crew_name="c", inputs={}, event_id="ck-1")) + listener.on_task_started( + None, SimpleNamespace(event_id="tk-1", task_name="t", agent_role="r", context={}, task=None) + ) + listener.on_agent_execution_started( + None, + SimpleNamespace( + event_id="ag-1", + agent=None, + agent_role="r", + task_prompt="p", + tools=[], + started_event_id="tk-1", + task_id=None, + ), + ) + listener.on_agent_execution_completed(None, SimpleNamespace(event_id="ag-2", started_event_id="ag-1", output="o")) + listener.on_task_completed(None, SimpleNamespace(event_id="tk-2", started_event_id="tk-1", output="o")) + listener.on_crew_kickoff_completed( + None, SimpleNamespace(event_id="ck-2", started_event_id="ck-1", output="o", total_tokens=1) + ) + + group_ids = {t["group_id"] for t in fake.task_messages} + assert len(fake.task_messages) == 3 + assert len(group_ids) == 1 + task_ids = {t["task_id"] for t in fake.task_messages} + assert len(task_ids) == 3 diff --git a/tests/agents/plugins/test_cross_plugin_linking.py b/tests/agents/plugins/test_cross_plugin_linking.py new file mode 100644 index 00000000..5238ecf5 --- /dev/null +++ b/tests/agents/plugins/test_cross_plugin_linking.py @@ -0,0 +1,160 @@ +"""Tests for cross-framework provenance linking through LangGraph. + +As documented in the top-level README ("Cross-framework provenance linking"), +a LangGraph run accepts ``_source_agent_id`` in the initial graph state to +link its provenance to a source task from another framework. The LangGraph +plugin stores the value as ``source_agent_id`` in ``custom_metadata`` of both +``langgraph_graph`` and ``langgraph_node`` records. + +Emission is captured in memory by replacing ``BaseInterceptor`` with a fake, +so no MQ, MongoDB, or network access is needed. +""" + +from __future__ import annotations + +import pytest + +import flowcept.flowceptor.adapters.base_interceptor as base_interceptor_module +from flowcept.agents.crewai.crewai_plugin import FlowceptCrewAIPlugin +from flowcept.agents.langgraph.langgraph_plugin import FlowceptLangGraphPlugin + +pytest.importorskip("langgraph") +pytest.importorskip("langchain_core") +pytest.importorskip("crewai") + + +class _CapturingInterceptor: + """In-memory stand-in for BaseInterceptor that records all emissions.""" + + instances: list = [] + + def __init__(self, plugin_key=None, kind=None): + """Record construction and initialize empty capture buffers.""" + self.kind = kind + self.telemetry_capture = None + self.stopped = False + self.task_messages: list[dict] = [] + self.workflow_messages: list = [] + type(self).instances.append(self) + + def start(self, bundle_exec_id, check_safe_stops=True): + """Mark the interceptor as started.""" + return self + + def stop(self, check_safe_stops=True): + """Mark the interceptor as stopped.""" + self.stopped = True + + def send_workflow_message(self, workflow_obj): + """Capture a WorkflowObject emission.""" + self.workflow_messages.append(workflow_obj) + + def intercept(self, task_dict): + """Capture a task message emission.""" + self.task_messages.append(task_dict) + + +@pytest.fixture() +def captured(monkeypatch): + """Patch BaseInterceptor with the capturing fake and return its instance list.""" + _CapturingInterceptor.instances = [] + monkeypatch.setattr(base_interceptor_module, "BaseInterceptor", _CapturingInterceptor) + return _CapturingInterceptor.instances + + +@pytest.fixture() +def langgraph_run(captured): + """Return a runner that invokes a local graph and returns the fake interceptor.""" + + def run(initial_state: dict): + plugin = FlowceptLangGraphPlugin(config={"workflow_name": "link-wf", "performance_tracking": False}) + plugin.start() + fake = captured[-1] + try: + graph = _build_graph() + graph.invoke(initial_state, config={"callbacks": [plugin.callback_handler]}) + finally: + plugin.stop() + return fake + + return run + + +def _build_graph(): + """Build a small local StateGraph of plain python-function nodes.""" + from typing import TypedDict + + from langgraph.graph import END, START, StateGraph + + class _State(TypedDict, total=False): + value: int + _source_agent_id: str + + def _add_one(state): + return {"value": state["value"] + 1} + + builder = StateGraph(_State) + builder.add_node("add_one", _add_one) + builder.add_edge(START, "add_one") + builder.add_edge("add_one", END) + return builder.compile() + + +def test_source_agent_id_is_stored_on_the_graph_task(langgraph_run): + """The langgraph_graph record carries custom_metadata.source_agent_id.""" + fake = langgraph_run({"value": 1, "_source_agent_id": "source-task-123"}) + graph_task = next(t for t in fake.task_messages if t["subtype"] == "langgraph_graph") + assert graph_task["custom_metadata"]["source_agent_id"] == "source-task-123" + # The raw linking key also travels in the recorded graph inputs. + assert graph_task["used"]["inputs"]["_source_agent_id"] == "source-task-123" + + +def test_source_agent_id_propagates_to_node_tasks(langgraph_run): + """Every langgraph_node record of the run carries the same source_agent_id.""" + fake = langgraph_run({"value": 1, "_source_agent_id": "source-task-123"}) + node_tasks = [t for t in fake.task_messages if t["subtype"] == "langgraph_node"] + assert node_tasks + for node_task in node_tasks: + assert node_task["custom_metadata"]["source_agent_id"] == "source-task-123" + + +def test_run_without_source_agent_id_has_no_linkage_field(langgraph_run): + """Without _source_agent_id in the state, no record carries source_agent_id.""" + fake = langgraph_run({"value": 1}) + assert fake.task_messages + for task in fake.task_messages: + assert "source_agent_id" not in task.get("custom_metadata", {}) + + +def test_langgraph_run_links_to_a_crewai_source_task(captured): + """A CrewAI task_id passed as _source_agent_id ends up on the LangGraph records.""" + from types import SimpleNamespace + + crew_plugin = FlowceptCrewAIPlugin(config={"workflow_name": "crew-src", "performance_tracking": False}) + crew_plugin.start() + crew_fake = captured[-1] + try: + listener = crew_plugin._listener_obj + listener.on_task_started( + None, SimpleNamespace(event_id="tk-1", task_name="research", agent_role="r", context={}, task=None) + ) + listener.on_task_completed(None, SimpleNamespace(event_id="tk-2", started_event_id="tk-1", output="data")) + source_task = next(t for t in crew_fake.task_messages if t["subtype"] == "crewai_task") + source_task_id = source_task["task_id"] + + lg_plugin = FlowceptLangGraphPlugin(config={"workflow_name": "lg-target", "performance_tracking": False}) + lg_plugin.start() + lg_fake = captured[-1] + try: + graph = _build_graph() + graph.invoke( + {"value": 1, "_source_agent_id": source_task_id}, + config={"callbacks": [lg_plugin.callback_handler]}, + ) + finally: + lg_plugin.stop() + finally: + crew_plugin.stop() + + graph_task = next(t for t in lg_fake.task_messages if t["subtype"] == "langgraph_graph") + assert graph_task["custom_metadata"]["source_agent_id"] == source_task_id diff --git a/tests/agents/plugins/test_langgraph_plugin.py b/tests/agents/plugins/test_langgraph_plugin.py new file mode 100644 index 00000000..ccf482c6 --- /dev/null +++ b/tests/agents/plugins/test_langgraph_plugin.py @@ -0,0 +1,357 @@ +"""Unit tests for the FlowCept LangGraph provenance plugin. + +Provenance emission is captured in memory by replacing ``BaseInterceptor`` +with a fake that records every workflow and task message, so no MQ, MongoDB, +or network access is needed. +""" + +from __future__ import annotations + +import uuid +from types import SimpleNamespace + +import pytest + +import flowcept.agents.langgraph.langgraph_plugin as lg_module +import flowcept.flowceptor.adapters.base_interceptor as base_interceptor_module +from flowcept.agents.langgraph.langgraph_plugin import FlowceptLangGraphPlugin + +pytest.importorskip("langgraph") +pytest.importorskip("langchain_core") + + +class _CapturingInterceptor: + """In-memory stand-in for BaseInterceptor that records all emissions.""" + + instances: list = [] + + def __init__(self, plugin_key=None, kind=None): + """Record construction and initialize empty capture buffers.""" + self.kind = kind + self.telemetry_capture = None + self.started = False + self.stopped = False + self.task_messages: list[dict] = [] + self.workflow_messages: list = [] + type(self).instances.append(self) + + def start(self, bundle_exec_id, check_safe_stops=True): + """Mark the interceptor as started.""" + self.started = True + return self + + def stop(self, check_safe_stops=True): + """Mark the interceptor as stopped.""" + self.stopped = True + + def send_workflow_message(self, workflow_obj): + """Capture a WorkflowObject emission.""" + self.workflow_messages.append(workflow_obj) + + def intercept(self, task_dict): + """Capture a task message emission.""" + self.task_messages.append(task_dict) + + +@pytest.fixture() +def captured(monkeypatch): + """Patch BaseInterceptor with the capturing fake and return its instance list.""" + _CapturingInterceptor.instances = [] + monkeypatch.setattr(base_interceptor_module, "BaseInterceptor", _CapturingInterceptor) + return _CapturingInterceptor.instances + + +@pytest.fixture() +def started_plugin(captured): + """Yield a started plugin plus the fake interceptor backing it.""" + plugin = FlowceptLangGraphPlugin(config={"workflow_name": "lg-test-wf", "performance_tracking": False}) + plugin.start() + assert captured, "plugin.start() did not build an interceptor" + yield plugin, captured[-1] + plugin.stop() + + +def _build_graph(failing_node=False): + """Build a small local StateGraph of plain python-function nodes.""" + from typing import TypedDict + + from langgraph.graph import END, START, StateGraph + + class _State(TypedDict, total=False): + value: int + _source_agent_id: str + + def _add_one(state): + return {"value": state["value"] + 1} + + def _double(state): + if failing_node: + raise ValueError("boom in node") + return {"value": state["value"] * 2} + + builder = StateGraph(_State) + builder.add_node("add_one", _add_one) + builder.add_node("double", _double) + builder.add_edge(START, "add_one") + builder.add_edge("add_one", "double") + builder.add_edge("double", END) + return builder.compile() + + +class _Generation: + """Minimal LangChain Generation stand-in.""" + + def __init__(self, text): + """Store the generated text.""" + self.text = text + + +class _LLMResult: + """Minimal LangChain LLMResult stand-in.""" + + def __init__(self, text, llm_output=None): + """Store one generation and optional llm_output metadata.""" + self.generations = [[_Generation(text)]] + self.llm_output = llm_output + + +# -- lifecycle ---------------------------------------------------------------- + + +def test_start_emits_top_level_workflow_message(started_plugin): + """start() sends one WorkflowObject carrying the configured workflow name.""" + _, fake = started_plugin + assert len(fake.workflow_messages) == 1 + wf = fake.workflow_messages[0] + assert wf.name == "lg-test-wf" + assert wf.workflow_id is not None + assert wf.campaign_id is not None + + +def test_start_respects_custom_campaign_id(captured): + """A campaign_id passed in config is used verbatim on the workflow message.""" + plugin = FlowceptLangGraphPlugin( + config={"workflow_name": "wf", "campaign_id": "camp-42", "performance_tracking": False} + ) + plugin.start() + try: + assert captured[-1].workflow_messages[0].campaign_id == "camp-42" + finally: + plugin.stop() + + +def test_stop_stops_the_underlying_interceptor(started_plugin): + """stop() flushes by stopping the wrapped interceptor exactly once.""" + plugin, fake = started_plugin + plugin.stop() + assert fake.stopped is True + plugin.stop() # second stop is a safe no-op + + +def test_disabled_plugin_emits_nothing(captured): + """enabled=False disables capture entirely and never builds an interceptor.""" + plugin = FlowceptLangGraphPlugin(config={"enabled": False}) + plugin.start() + assert captured == [] + with pytest.raises(RuntimeError): + _ = plugin.callback_handler + plugin.stop() + + +def test_callback_handler_raises_before_start(captured): + """Accessing callback_handler before start() raises RuntimeError.""" + plugin = FlowceptLangGraphPlugin(config={"performance_tracking": False}) + with pytest.raises(RuntimeError): + _ = plugin.callback_handler + + +def test_context_manager_starts_and_stops(captured): + """The plugin works as a context manager, starting on enter and stopping on exit.""" + with FlowceptLangGraphPlugin(config={"workflow_name": "ctx-wf", "performance_tracking": False}) as plugin: + assert plugin._started is True + fake = captured[-1] + assert plugin._started is False + assert fake.stopped is True + + +# -- graph runs --------------------------------------------------------------- + + +def test_graph_invoke_emits_graph_and_node_tasks(started_plugin): + """A local graph run yields one langgraph_graph task and one task per node.""" + plugin, fake = started_plugin + graph = _build_graph() + result = graph.invoke({"value": 3}, config={"callbacks": [plugin.callback_handler]}) + assert result["value"] == 8 + + by_subtype = {} + for task in fake.task_messages: + by_subtype.setdefault(task["subtype"], []).append(task) + graph_task = by_subtype["langgraph_graph"][0] + node_names = {t["activity_id"] for t in by_subtype["langgraph_node"]} + assert {"add_one", "double"}.issubset(node_names) + assert graph_task["status"] == "FINISHED" + assert graph_task["used"]["inputs"]["value"] == 3 + assert graph_task["generated"]["outputs"]["value"] == 8 + for task in fake.task_messages: + assert task["workflow_id"] == fake.workflow_messages[0].workflow_id + assert task["campaign_id"] == fake.workflow_messages[0].campaign_id + assert task["status"] == "FINISHED" + + +def test_graph_invocation_emits_sub_workflow_linked_to_parent(started_plugin): + """Each graph.invoke emits a sub-WorkflowObject pointing at the top workflow.""" + plugin, fake = started_plugin + graph = _build_graph() + graph.invoke({"value": 1}, config={"callbacks": [plugin.callback_handler]}) + + top_wf, sub_wf = fake.workflow_messages[0], fake.workflow_messages[1] + assert sub_wf.parent_workflow_id == top_wf.workflow_id + assert sub_wf.custom_metadata["graph_name"] == sub_wf.name + graph_task = next(t for t in fake.task_messages if t["subtype"] == "langgraph_graph") + assert sub_wf.custom_metadata["group_id"] == graph_task["group_id"] + + +def test_node_tasks_share_group_id_and_link_to_graph_task(started_plugin): + """All tasks of one invocation share a group_id; nodes parent to the graph task.""" + plugin, fake = started_plugin + graph = _build_graph() + graph.invoke({"value": 1}, config={"callbacks": [plugin.callback_handler]}) + + graph_task = next(t for t in fake.task_messages if t["subtype"] == "langgraph_graph") + node_tasks = [t for t in fake.task_messages if t["subtype"] == "langgraph_node"] + assert node_tasks + task_ids = {t["task_id"] for t in fake.task_messages} + assert len(task_ids) == len(fake.task_messages) # unique task ids + for node_task in node_tasks: + assert node_task["group_id"] == graph_task["group_id"] + assert node_task["parent_task_id"] == graph_task["task_id"] + + +def test_node_error_is_recorded_as_failed_task(started_plugin): + """A raising node produces ERROR-status tasks carrying the exception text.""" + plugin, fake = started_plugin + graph = _build_graph(failing_node=True) + with pytest.raises(ValueError, match="boom in node"): + graph.invoke({"value": 1}, config={"callbacks": [plugin.callback_handler]}) + + failed = [t for t in fake.task_messages if t["status"] == "ERROR"] + assert failed + assert any("boom in node" in t.get("stderr", "") for t in failed) + + +# -- LLM and tool callback events ---------------------------------------------- + + +def test_llm_events_emit_llm_call_task(started_plugin): + """on_llm_start/on_llm_end produce one llm_call task with prompts and tokens.""" + plugin, fake = started_plugin + handler = plugin.callback_handler + parent_id, run_id = uuid.uuid4(), uuid.uuid4() + handler.on_chain_start(None, {"q": "hi"}, run_id=parent_id) + handler.on_llm_start({"kwargs": {"model": "fake-model"}}, ["what is 2+2"], run_id=run_id, parent_run_id=parent_id) + result = _LLMResult("4", {"token_usage": {"prompt_tokens": 12, "completion_tokens": 1, "total_tokens": 13}}) + handler.on_llm_end(result, run_id=run_id, parent_run_id=parent_id) + handler.on_chain_end({"a": "4"}, run_id=parent_id) + + llm_task = next(t for t in fake.task_messages if t["subtype"] == "llm_call") + graph_task = next(t for t in fake.task_messages if t["subtype"] == "langgraph_graph") + assert llm_task["activity_id"] == "fake-model" + assert llm_task["used"]["prompts"] == ["what is 2+2"] + assert llm_task["generated"]["text"] == "4" + assert llm_task["generated"]["total_tokens"] == 13 + assert llm_task["status"] == "FINISHED" + assert llm_task["parent_task_id"] == graph_task["task_id"] + assert llm_task["group_id"] == graph_task["group_id"] + + +def test_llm_error_emits_error_task(started_plugin): + """on_llm_error records the llm_call task as ERROR with stderr set.""" + plugin, fake = started_plugin + handler = plugin.callback_handler + run_id = uuid.uuid4() + handler.on_llm_start({"kwargs": {"model": "m"}}, ["p"], run_id=run_id) + handler.on_llm_error(RuntimeError("rate limited"), run_id=run_id) + + llm_task = next(t for t in fake.task_messages if t["subtype"] == "llm_call") + assert llm_task["status"] == "ERROR" + assert "rate limited" in llm_task["stderr"] + + +def test_chat_model_events_emit_llm_call_task(started_plugin): + """on_chat_model_start serializes message contents into the llm_call task.""" + plugin, fake = started_plugin + handler = plugin.callback_handler + run_id = uuid.uuid4() + messages = [[SimpleNamespace(content="hello"), SimpleNamespace(content="world")]] + handler.on_chat_model_start({"kwargs": {"model": "chat-model"}}, messages, run_id=run_id) + handler.on_llm_end(_LLMResult("hi"), run_id=run_id) + + llm_task = next(t for t in fake.task_messages if t["subtype"] == "llm_call") + assert llm_task["used"]["messages"] == [["hello", "world"]] + assert llm_task["used"]["model"] == "chat-model" + assert llm_task["generated"]["text"] == "hi" + + +def test_tool_events_emit_tool_call_task(started_plugin): + """on_tool_start/on_tool_end produce one tool_call task with input and output.""" + plugin, fake = started_plugin + handler = plugin.callback_handler + parent_id, run_id = uuid.uuid4(), uuid.uuid4() + handler.on_chain_start(None, {}, run_id=parent_id) + handler.on_tool_start({"name": "get_weather"}, '{"city": "Paris"}', run_id=run_id, parent_run_id=parent_id) + handler.on_tool_end("18C", run_id=run_id, parent_run_id=parent_id) + handler.on_chain_end({}, run_id=parent_id) + + tool_task = next(t for t in fake.task_messages if t["subtype"] == "tool_call") + graph_task = next(t for t in fake.task_messages if t["subtype"] == "langgraph_graph") + assert tool_task["activity_id"] == "get_weather" + assert tool_task["used"]["input"] == '{"city": "Paris"}' + assert tool_task["generated"]["output"] == "18C" + assert tool_task["status"] == "FINISHED" + assert tool_task["parent_task_id"] == graph_task["task_id"] + + +def test_tool_error_emits_error_task(started_plugin): + """on_tool_error records the tool_call task as ERROR with stderr set.""" + plugin, fake = started_plugin + handler = plugin.callback_handler + run_id = uuid.uuid4() + handler.on_tool_start({"name": "deploy"}, "", run_id=run_id) + handler.on_tool_error(RuntimeError("denied"), run_id=run_id) + + tool_task = next(t for t in fake.task_messages if t["subtype"] == "tool_call") + assert tool_task["status"] == "ERROR" + assert "denied" in tool_task["stderr"] + + +# -- record_llm_call public API ------------------------------------------------- + + +def test_record_llm_call_routes_through_active_interceptor(started_plugin): + """record_llm_call emits an llm_call task via the module-level interceptor.""" + _, fake = started_plugin + lg_module.record_llm_call( + { + "type": "chat_completion", + "model": "gpt-test", + "text": "hello", + "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + } + ) + llm_task = next(t for t in fake.task_messages if t["subtype"] == "llm_call") + assert llm_task["activity_id"] == "gpt-test" + assert llm_task["generated"]["text"] == "hello" + assert llm_task["used"]["model"] == "gpt-test" + assert llm_task["status"] == "FINISHED" + + +def test_record_llm_call_is_a_noop_when_plugin_stopped(captured): + """After stop(), record_llm_call does not emit anything.""" + plugin = FlowceptLangGraphPlugin(config={"performance_tracking": False}) + plugin.start() + fake = captured[-1] + plugin.stop() + emitted_before = len(fake.task_messages) + lg_module.record_llm_call({"type": "chat_completion", "model": "m", "text": "t", "usage": {}}) + assert len(fake.task_messages) == emitted_before diff --git a/tests/agents/prov_analysis/__init__.py b/tests/agents/prov_analysis/__init__.py new file mode 100644 index 00000000..50ef3996 --- /dev/null +++ b/tests/agents/prov_analysis/__init__.py @@ -0,0 +1 @@ +"""Tests for the agentic provenance analysis package.""" diff --git a/tests/agents/prov_analysis/test_core.py b/tests/agents/prov_analysis/test_core.py new file mode 100644 index 00000000..a3cd901f --- /dev/null +++ b/tests/agents/prov_analysis/test_core.py @@ -0,0 +1,450 @@ +"""Unit tests for the provenance analysis core. + +Synthetic records cover both shapes the core must understand: the +harness-buffer shape (``agents/harness/prov.py``: ``type`` on every record, +``custom_metadata.llm_usage``) and the framework-plugin shape (LangGraph-style +task dicts without a ``type`` key, token counts in ``generated``). No MQ, +MongoDB, network, or LLM access is needed. +""" + +from __future__ import annotations + +import json + +import pytest + +from flowcept.agents.prov_analysis import core + + +# -- synthetic records --------------------------------------------------------- + + +def harness_records() -> list[dict]: + """One harness session with a subagent, mirroring recorder.py output.""" + return [ + { + "type": "workflow", + "workflow_id": "wf-h", + "campaign_id": "camp-1", + "name": "claude_code session", + "subtype": "agent_session", + "agent_id": "agent-main", + "status": "FINISHED", + "started_at": 1000.0, + "ended_at": 1060.0, + "custom_metadata": {"harness": "claude_code"}, + }, + { + "type": "workflow", + "workflow_id": "wf-sub", + "parent_workflow_id": "wf-h", + "name": "subagent:Explore", + "subtype": "subagent_session", + "status": "FINISHED", + "started_at": 1020.0, + "ended_at": 1030.0, + }, + {"type": "agent", "agent_id": "agent-main", "name": "claude"}, + { + "type": "task", + "task_id": "turn-1", + "workflow_id": "wf-h", + "campaign_id": "camp-1", + "activity_id": "agent_turn", + "subtype": "ai_model_invocation", + "agent_id": "agent-main", + "status": "FINISHED", + "started_at": 1001.0, + "ended_at": 1050.0, + "used": {"prompt": "fix the flake"}, + "generated": {"response": "done"}, + "custom_metadata": { + "granularity": "turn", + "harness": "claude_code", + "llm_usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150}, + }, + }, + { + "type": "task", + "task_id": "tool-1", + "workflow_id": "wf-h", + "activity_id": "Bash", + "subtype": "agent_tool", + "agent_id": "agent-main", + "parent_task_id": "turn-1", + "status": "FINISHED", + "started_at": 1002.0, + "ended_at": 1010.0, + "used": {"command": "pytest -q"}, + "custom_metadata": {"harness": "claude_code", "tool_name": "Bash"}, + }, + { + "type": "task", + "task_id": "tool-2", + "workflow_id": "wf-h", + "activity_id": "Bash", + "subtype": "agent_tool", + "agent_id": "agent-main", + "parent_task_id": "turn-1", + "status": "ERROR", + "started_at": 1011.0, + "ended_at": 1012.0, + "stderr": "1 error found", + "custom_metadata": {"harness": "claude_code", "tool_name": "Bash"}, + }, + { + "type": "task", + "task_id": "tool-3", + "workflow_id": "wf-sub", + "activity_id": "Grep", + "subtype": "agent_tool", + "agent_id": "agent-sub", + "parent_task_id": "tool-1", + "status": "FINISHED", + "started_at": 1021.0, + "ended_at": 1023.0, + }, + { + "type": "task", + "task_id": "llm-1", + "workflow_id": "wf-h", + "activity_id": "llm_interaction", + "subtype": "ai_model_invocation", + "agent_id": "agent-main", + "parent_task_id": "turn-1", + "status": "FINISHED", + "started_at": 1030.0, + "ended_at": 1031.0, + "custom_metadata": {"granularity": "call", "llm_usage": {"total_tokens": 25}}, + }, + ] + + +def framework_records() -> list[dict]: + """LangGraph-plugin-shaped task dicts (no ``type`` key) with a cross link.""" + return [ + { + "task_id": "lg-graph-1", + "workflow_id": "wf-lg", + "activity_id": "LangGraph", + "subtype": "langgraph_graph", + "status": "FINISHED", + "started_at": 2000.0, + "ended_at": 2005.0, + "used": {"inputs": {"value": 1, "_source_agent_id": "tool-1"}}, + "custom_metadata": {"graph_name": "LangGraph", "source_agent_id": "tool-1"}, + }, + { + "task_id": "lg-node-1", + "workflow_id": "wf-lg", + "activity_id": "add_one", + "subtype": "langgraph_node", + "parent_task_id": "lg-graph-1", + "status": "FINISHED", + "started_at": 2001.0, + "ended_at": 2003.0, + "custom_metadata": {"node_name": "add_one", "source_agent_id": "tool-1"}, + }, + { + "task_id": "lg-llm-1", + "workflow_id": "wf-lg", + "activity_id": "gpt-test", + "subtype": "llm_call", + "parent_task_id": "lg-node-1", + "status": "FINISHED", + "started_at": 2001.5, + "ended_at": 2002.5, + "generated": {"text": "hi", "prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + { + "task_id": "lg-tool-1", + "workflow_id": "wf-lg", + "activity_id": "search", + "subtype": "tool_call", + "parent_task_id": "lg-node-1", + "status": "ERROR", + "started_at": 2003.0, + "ended_at": 2004.0, + "stderr": "boom", + }, + ] + + +# -- load_records -------------------------------------------------------------- + + +def test_load_records_from_jsonl(tmp_path): + """load_records reads dict lines from JSONL and skips garbage lines.""" + path = tmp_path / "buffer.jsonl" + lines = [json.dumps(r) for r in harness_records()] + ["", "not json", '["not", "a", "dict"]'] + path.write_text("\n".join(lines), encoding="utf-8") + records = core.load_records(jsonl_path=str(path)) + assert len(records) == len(harness_records()) + assert records[0]["workflow_id"] == "wf-h" + + +def test_load_records_passes_records_through(): + """load_records returns a copy of an explicitly passed record list.""" + given = harness_records() + out = core.load_records(records=given) + assert out == given + assert out is not given, "must return a copy, not the caller's list" + + +def test_load_records_with_no_source_is_empty(): + """load_records with no source returns an empty list.""" + assert core.load_records() == [] + + +# -- summarize_execution ------------------------------------------------------- + + +def test_summarize_execution_counts_and_bounds(): + """Summary reports counts, statuses, bounds, campaigns, and agents.""" + summary = core.summarize_execution(harness_records()) + assert summary["n_workflows"] == 2 + assert summary["n_tasks"] == 5 + assert summary["tasks_by_subtype"] == {"ai_model_invocation": 2, "agent_tool": 3} + assert summary["tasks_by_activity"]["Bash"] == 2 + assert summary["status_counts"] == {"FINISHED": 4, "ERROR": 1} + assert summary["started_at"] == 1001.0 + assert summary["ended_at"] == 1050.0 + assert summary["total_elapsed_seconds"] == pytest.approx(49.0) + assert summary["campaigns"] == ["camp-1"] + assert "agent-main" in summary["agents"] + assert summary["agent_names"] == ["claude"] + + +def test_summarize_execution_token_usage_from_both_shapes(): + """Token totals combine harness llm_usage and plugin generated fields.""" + summary = core.summarize_execution(harness_records() + framework_records()) + totals = summary["token_usage"]["totals"] + # 100+10 prompt, 50+5 completion, 150+25+15 total + assert totals["prompt_tokens"] == 110 + assert totals["completion_tokens"] == 55 + assert totals["total_tokens"] == 190 + assert summary["token_usage"]["n_tasks_with_usage"] == 3 + + +def test_summarize_execution_filters_by_workflow_id(): + """workflow_id restricts the summary to one workflow's records.""" + summary = core.summarize_execution(harness_records(), workflow_id="wf-sub") + assert summary["n_tasks"] == 1 + assert summary["tasks_by_activity"] == {"Grep": 1} + assert summary["n_workflows"] == 1 + + +def test_summarize_execution_handles_framework_tasks_without_type(): + """Plugin task dicts without a type key still count as tasks.""" + summary = core.summarize_execution(framework_records()) + assert summary["n_tasks"] == 4 + assert summary["n_workflows"] == 0 + assert summary["tasks_by_subtype"]["langgraph_node"] == 1 + + +def test_summarize_execution_empty_input(): + """An empty record list yields a zeroed summary.""" + summary = core.summarize_execution([]) + assert summary["n_records"] == 0 + assert summary["n_tasks"] == 0 + assert summary["total_elapsed_seconds"] is None + assert summary["token_usage"]["totals"] == {} + assert summary["activities"] == [] + + +# -- analyze_errors ------------------------------------------------------------ + + +def test_analyze_errors_groups_by_activity_with_excerpts(): + """Failures group by activity with rates and stderr excerpts.""" + errors = core.analyze_errors(harness_records()) + assert errors["n_failed"] == 1 + assert errors["overall_error_rate"] == pytest.approx(0.2) + entry = errors["by_activity"]["Bash"] + assert entry["n_failed"] == 1 + assert entry["n_total"] == 2 + assert entry["error_rate"] == pytest.approx(0.5) + assert entry["excerpts"] == ["1 error found"] + + +def test_analyze_errors_failure_time_bounds(): + """First/last failure times span both record shapes.""" + errors = core.analyze_errors(harness_records() + framework_records()) + assert errors["n_failed"] == 2 + assert errors["first_failure_at"] == 1011.0 + assert errors["last_failure_at"] == 2003.0 + assert errors["first_failure_at_utc"] is not None + + +def test_analyze_errors_without_failures(): + """No failed tasks yields empty groupings and a zero rate.""" + ok_only = [r for r in harness_records() if r.get("status") != "ERROR"] + errors = core.analyze_errors(ok_only) + assert errors["n_failed"] == 0 + assert errors["by_activity"] == {} + assert errors["first_failure_at"] is None + assert errors["overall_error_rate"] == 0.0 + + +def test_analyze_errors_empty_input(): + """An empty record list yields no error rate.""" + errors = core.analyze_errors([]) + assert errors["n_tasks"] == 0 + assert errors["overall_error_rate"] is None + + +# -- analyze_agent_behavior ------------------------------------------------------ + + +def test_agent_behavior_per_agent_counts(): + """Per-agent turns, tool calls, LLM calls, tokens, and durations.""" + behavior = core.analyze_agent_behavior(harness_records()) + main = behavior["agents"]["agent-main"] + assert main["turns"] == 1 + assert main["llm_calls"] == 1 # granularity=call record + assert main["tool_calls"] == 2 + assert main["tool_calls_by_tool"] == {"Bash": 2} + assert main["n_errors"] == 1 + assert main["token_usage"]["total_tokens"] == 175 + assert main["max_task_seconds"] == pytest.approx(49.0) + sub = behavior["agents"]["agent-sub"] + assert sub["tool_calls_by_tool"] == {"Grep": 1} + + +def test_agent_behavior_sessions_and_subagents(): + """Session workflows and subagent counts are reported.""" + behavior = core.analyze_agent_behavior(harness_records()) + assert behavior["n_subagent_sessions"] == 1 + assert len(behavior["sessions"]) == 1 + session = behavior["sessions"][0] + assert session["workflow_id"] == "wf-h" + assert session["n_subagents"] == 1 + assert session["elapsed_seconds"] == pytest.approx(60.0) + + +def test_agent_behavior_framework_shape_falls_back_to_workflow_key(): + """Plugin records without agent_id key by workflow_id.""" + behavior = core.analyze_agent_behavior(framework_records()) + entry = behavior["agents"]["wf-lg"] + assert entry["llm_calls"] == 1 + assert entry["tool_calls_by_tool"] == {"search": 1} + assert entry["token_usage"]["total_tokens"] == 15 + + +def test_agent_behavior_empty_input(): + """An empty record list yields an empty behavior profile.""" + behavior = core.analyze_agent_behavior([]) + assert behavior == {"agents": {}, "sessions": [], "n_subagent_sessions": 0} + + +# -- find_slowest_tasks ---------------------------------------------------------- + + +def test_find_slowest_orders_and_limits(): + """Slowest tasks come back longest first, capped by limit.""" + rows = core.find_slowest_tasks(harness_records(), limit=3) + assert [r["task_id"] for r in rows] == ["turn-1", "tool-1", "tool-3"] + assert rows[0]["elapsed_seconds"] == pytest.approx(49.0) + assert rows[0]["status"] == "FINISHED" + + +def test_find_slowest_reports_parent_depth(): + """Parent-chain depth follows parent_task_id links.""" + rows = core.find_slowest_tasks(harness_records(), limit=10) + by_id = {r["task_id"]: r for r in rows} + assert by_id["turn-1"]["parent_depth"] == 0 + assert by_id["tool-1"]["parent_depth"] == 1 + assert by_id["tool-3"]["parent_depth"] == 2 + + +def test_find_slowest_skips_tasks_without_timing(): + """Tasks without timing are excluded; empty input yields [].""" + records = harness_records() + [{"type": "task", "task_id": "no-time", "activity_id": "X"}] + rows = core.find_slowest_tasks(records, limit=100) + assert all(r["task_id"] != "no-time" for r in rows) + assert core.find_slowest_tasks([], limit=5) == [] + + +# -- cross_framework_links -------------------------------------------------------- + + +def test_cross_framework_links_builds_edges(): + """Edges are built from source_agent_id pointers across shapes.""" + result = core.cross_framework_links(harness_records() + framework_records()) + assert result["n_links"] == 2 + targets = {link["target_task_id"]: link for link in result["links"]} + assert set(targets) == {"lg-graph-1", "lg-node-1"} + graph_link = targets["lg-graph-1"] + assert graph_link["source_task_id"] == "tool-1" + assert graph_link["target_workflow_id"] == "wf-lg" + # Source is a harness (claude_code) task; target is a langgraph task. + assert graph_link["frameworks"] == ["claude_code", "langgraph"] + + +def test_cross_framework_links_from_used_inputs_only(): + """The raw _source_agent_id inside used.inputs also links.""" + records = [ + { + "task_id": "t1", + "workflow_id": "w1", + "subtype": "langgraph_graph", + "used": {"inputs": {"_source_agent_id": "external-1"}}, + } + ] + result = core.cross_framework_links(records) + assert result["n_links"] == 1 + assert result["links"][0]["source_task_id"] == "external-1" + assert result["n_unlinked_tasks"] == 0 + + +def test_cross_framework_links_counts_unlinked(): + """Sessions without pointers report zero links and all tasks unlinked.""" + result = core.cross_framework_links(harness_records()) + assert result["n_links"] == 0 + assert result["n_unlinked_tasks"] == 5 + assert core.cross_framework_links([]) == { + "links": [], + "n_links": 0, + "n_unlinked_tasks": 0, + "frameworks_seen": [], + } + + +# -- compare_executions ------------------------------------------------------------ + + +def test_compare_executions_deltas(): + """Per-activity count, duration, and error-rate deltas are computed.""" + records_a = harness_records() + records_b = [dict(r) for r in harness_records()] + # Make run B's failing Bash task succeed and take longer. + for record in records_b: + if record.get("task_id") == "tool-2": + record["status"] = "FINISHED" + record["ended_at"] = 1021.0 + result = core.compare_executions(records_a, records_b) + bash = result["activities"]["Bash"] + assert bash["count_a"] == bash["count_b"] == 2 + assert bash["count_delta"] == 0 + assert bash["error_rate_a"] == pytest.approx(0.5) + assert bash["error_rate_b"] == 0.0 + assert bash["error_rate_delta"] == pytest.approx(-0.5) + assert bash["elapsed_avg_delta"] == pytest.approx(4.5) + assert result["totals"]["n_tasks_a"] == 5 + + +def test_compare_executions_disjoint_activities(): + """Activities present in only one run are listed separately.""" + result = core.compare_executions(harness_records(), framework_records()) + assert "Bash" in result["only_in_a"] + assert "add_one" in result["only_in_b"] + assert result["activities"]["Bash"]["count_b"] == 0 + assert result["activities"]["Bash"]["elapsed_avg_b"] is None + + +def test_compare_executions_empty_inputs(): + """Comparing empty runs yields empty activities and null deltas.""" + result = core.compare_executions([], []) + assert result["activities"] == {} + assert result["totals"]["n_tasks_a"] == 0 + assert result["totals"]["total_elapsed_delta"] is None diff --git a/tests/agents/prov_analysis/test_registration.py b/tests/agents/prov_analysis/test_registration.py new file mode 100644 index 00000000..4a66dde6 --- /dev/null +++ b/tests/agents/prov_analysis/test_registration.py @@ -0,0 +1,110 @@ +"""Registration and offline behavior of the Flowcept agent analysis MCP tools. + +Importing ``analysis_mcp_tools`` must register every tool on ``mcp_flowcept``. +Skips gracefully when the MCP/agent stack cannot be imported in this +environment (heavy config or missing extras); no MQ, MongoDB, network, or LLM +keys are needed otherwise. +""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("mcp") +pytest.importorskip("pandas") + +try: + import flowcept.agents.mcp.mcp_tools.analysis_mcp_tools as analysis_mcp_tools + from flowcept.agents.mcp.context_manager import ctx_manager, mcp_flowcept +except Exception as exc: # pragma: no cover - environment-dependent + pytest.skip(f"agent MCP stack unavailable: {exc}", allow_module_level=True) + +EXPECTED_TOOLS = { + "df_summarize_execution", + "df_analyze_errors", + "df_agent_behavior", + "df_find_slowest", + "df_cross_framework_links", + "db_summarize_execution", + "db_analyze_errors", + "db_agent_behavior", + "db_find_slowest", + "db_cross_framework_links", + "compare_executions", +} + + +def test_analysis_tools_are_registered(): + """The four analysis tools are registered on the harness server.""" + names = {t.name for t in mcp_flowcept._tool_manager.list_tools()} + assert EXPECTED_TOOLS <= names + + +def test_analysis_tools_have_descriptions(): + """Every analysis tool carries a description for the model.""" + by_name = {t.name: t for t in mcp_flowcept._tool_manager.list_tools()} + for name in EXPECTED_TOOLS: + assert by_name[name].description, f"{name} has no description for the model to read" + + +def test_df_tools_report_empty_context(): + """DF tools return 404 when no records are loaded.""" + ctx_manager.context.reset_context() + result = analysis_mcp_tools.df_summarize_execution() + assert result.code == 404 + + +def test_df_tools_analyze_loaded_context_records(): + """DF tools analyze the raw records held in the agent context.""" + ctx_manager.context.reset_context() + ctx_manager.context.workflow_msg_obj = { + "type": "workflow", + "workflow_id": "wf-1", + "status": "FINISHED", + } + ctx_manager.context.tasks = [ + { + "type": "task", + "task_id": "t1", + "workflow_id": "wf-1", + "activity_id": "train", + "subtype": "agent_tool", + "status": "FINISHED", + "started_at": 1.0, + "ended_at": 3.0, + }, + { + "type": "task", + "task_id": "t2", + "workflow_id": "wf-1", + "activity_id": "train", + "subtype": "agent_tool", + "status": "ERROR", + "stderr": "exploded", + "started_at": 3.0, + "ended_at": 4.0, + }, + ] + try: + summary = analysis_mcp_tools.df_summarize_execution() + assert summary.code == 301 + assert summary.result["n_tasks"] == 2 + assert summary.result["status_counts"]["ERROR"] == 1 + + errors = analysis_mcp_tools.df_analyze_errors() + assert errors.code == 301 + assert errors.result["by_activity"]["train"]["excerpts"] == ["exploded"] + + slowest = analysis_mcp_tools.df_find_slowest(limit=1) + assert slowest.code == 301 + assert slowest.result["tasks"][0]["task_id"] == "t1" + + links = analysis_mcp_tools.df_cross_framework_links() + assert links.code == 301 + assert links.result["n_links"] == 0 + + behavior = analysis_mcp_tools.df_agent_behavior() + assert behavior.code == 301 + assert behavior.result["agents"]["wf-1"]["tool_calls"] == 2 + finally: + ctx_manager.context.reset_context() diff --git a/tests/harness/__init__.py b/tests/harness/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/harness/conftest.py b/tests/harness/conftest.py new file mode 100644 index 00000000..fe230111 --- /dev/null +++ b/tests/harness/conftest.py @@ -0,0 +1,34 @@ +"""Shared fixtures. + +Every test gets its own capture home so that state files and buffers from one +test can never be seen by another. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from flowcept.agents.harness.config import Config + + +@pytest.fixture +def config(tmp_path: Path) -> Config: + return Config(home=tmp_path / "home", debug=True) + + +@pytest.fixture +def buffer_records(config: Config): + """Return a reader for every record written to any buffer.""" + + def read() -> list[dict]: + records: list[dict] = [] + for path in sorted(config.buffers_dir.glob("*.jsonl")): + for line in path.read_text(encoding="utf-8").splitlines(): + if line.strip(): + records.append(json.loads(line)) + return records + + return read diff --git a/tests/harness/test_analysis_tools.py b/tests/harness/test_analysis_tools.py new file mode 100644 index 00000000..0b719237 --- /dev/null +++ b/tests/harness/test_analysis_tools.py @@ -0,0 +1,250 @@ +"""Tests for the provenance analysis tools on the harness MCP server and CLI. + +The MCP tools are exercised through the server's own registry (like +``test_mcp_server.py``) against a captured synthetic session; the CLI +``analyze`` subcommand is driven like ``test_cli.py`` does. Everything runs +offline: no MQ, MongoDB, network, or LLM keys. +""" + +from __future__ import annotations + +import json + +import pytest + +from flowcept.agents.harness import cli, ids + +from .test_claude_code import SESSION, fire +from .test_mcp_server import call + +pytest.importorskip("mcp") + +from flowcept.agents.harness import mcp_server # noqa: E402 + + +def record_session(config): + """Capture one synthetic session with a failure and a subagent.""" + fire(config, "SessionStart", source="startup", model="claude-opus-5") + fire(config, "UserPromptSubmit", prompt="fix the flake", prompt_id="p1") + fire(config, "PreToolUse", tool_name="Bash", tool_use_id="t1", tool_input={"command": "pytest -q"}) + fire(config, "PostToolUse", tool_name="Bash", tool_use_id="t1", tool_response={"exit_code": 0}) + fire(config, "PreToolUse", tool_name="Bash", tool_use_id="t2", tool_input={"command": "ruff check"}) + fire(config, "PostToolUseFailure", tool_name="Bash", tool_use_id="t2", error="1 error found") + fire(config, "SubagentStart", agent_id="a1", agent_type="Explore") + fire(config, "PreToolUse", tool_name="Grep", tool_use_id="t3", tool_input={"pattern": "flaky"}, agent_id="a1") + fire(config, "PostToolUse", tool_name="Grep", tool_use_id="t3", tool_response={"matches": 2}, agent_id="a1") + fire(config, "SubagentStop", agent_id="a1", agent_type="Explore") + fire(config, "Stop", last_assistant_message="Fixed the race.") + fire(config, "SessionEnd", reason="clear") + + +@pytest.fixture +def server(config): + """Build an MCP server over one captured synthetic session.""" + record_session(config) + return mcp_server.build_server(config) + + +def test_analysis_tools_are_registered(server): + """The four analysis tools are registered on the harness server.""" + names = {t.name for t in server._tool_manager.list_tools()} + assert {"analyze_session", "analyze_errors", "find_slowest", "cross_links"} <= names + + +def test_analysis_tools_have_descriptions(server): + """Every analysis tool carries a description for the model.""" + for name in ("analyze_session", "analyze_errors", "find_slowest", "cross_links"): + tool = server._tool_manager.get_tool(name) + assert tool.description + + +def test_analyze_session_summary_and_behavior(server): + """analyze_session returns a summary plus agent behavior.""" + result = call(server, "analyze_session") + summary = result["summary"] + assert summary["n_workflows"] == 2 + assert summary["tasks_by_activity"]["Bash"] == 2 + assert summary["status_counts"]["ERROR"] == 1 + behavior = result["agent_behavior"] + assert behavior["n_subagent_sessions"] == 1 + assert behavior["sessions"][0]["n_subagents"] == 1 + tools_seen = {tool for entry in behavior["agents"].values() for tool in entry["tool_calls_by_tool"]} + assert {"Bash", "Grep"} <= tools_seen + + +def test_analyze_session_reports_a_miss(server): + """An unknown session id returns an error payload.""" + assert "error" in call(server, "analyze_session", session="nope") + + +def test_analyze_errors_tool(server): + """analyze_errors reports per-activity failures with excerpts.""" + errors = call(server, "analyze_errors") + assert errors["n_failed"] == 1 + assert errors["by_activity"]["Bash"]["excerpts"] == ["1 error found"] + assert errors["by_activity"]["Bash"]["error_rate"] == pytest.approx(0.5) + + +def test_analyze_errors_reports_a_miss(server): + """An unknown session id returns an error payload.""" + assert "error" in call(server, "analyze_errors", session="nope") + + +def test_find_slowest_tool(server): + """find_slowest returns ordered rows with the expected fields.""" + rows = call(server, "find_slowest", limit=2) + assert len(rows) == 2 + assert rows[0]["elapsed_seconds"] >= rows[1]["elapsed_seconds"] + assert {"task_id", "activity_id", "status", "parent_depth"} <= set(rows[0]) + + +def test_cross_links_tool_finds_a_planted_link(config): + """A planted plugin record linking a harness task is surfaced.""" + record_session(config) + # Plant a framework-plugin task linking back to a harness tool task. + buffer = next(iter(config.buffers_dir.glob("*.jsonl"))) + records = [json.loads(line) for line in buffer.read_text().splitlines() if line.strip()] + tool_task = next(r for r in records if r.get("subtype") == "agent_tool") + linked = { + "task_id": "lg-1", + "workflow_id": "wf-lg", + "activity_id": "graph", + "subtype": "langgraph_graph", + "status": "FINISHED", + "custom_metadata": {"source_agent_id": tool_task["task_id"]}, + } + with buffer.open("a", encoding="utf-8") as fh: + fh.write(json.dumps(linked) + "\n") + + server = mcp_server.build_server(config) + result = call(server, "cross_links") + assert result["n_links"] == 1 + link = result["links"][0] + assert link["source_task_id"] == tool_task["task_id"] + assert link["target_task_id"] == "lg-1" + assert "langgraph" in link["frameworks"] + + +def test_cross_links_without_links(server): + """A plain session reports zero links and unlinked tasks.""" + result = call(server, "cross_links") + assert result["n_links"] == 0 + assert result["n_unlinked_tasks"] > 0 + + +# -- CLI ``analyze`` subcommand ------------------------------------------------- + + +@pytest.fixture +def run(config, monkeypatch, capsys): + """Invoke the CLI against the test's capture home.""" + monkeypatch.setenv("FLOWCEPT_HARNESS_HOME", str(config.home)) + + def _run(*argv: str): + code = cli.main(list(argv)) + captured = capsys.readouterr() + return code, captured.out + captured.err + + return _run + + +def test_cli_analyze_summary(run, config): + """`analyze` prints counts, statuses, subagents, and tool usage.""" + record_session(config) + code, out = run("analyze") + assert code == cli.OK + assert "tasks: 4" in out + assert "ERROR=1" in out + assert "subagents=1" in out + assert "Bash=2" in out + + +def test_cli_analyze_errors(run, config): + """`analyze --errors` prints failure counts and excerpts.""" + record_session(config) + code, out = run("analyze", "--errors") + assert code == cli.OK + assert "failed tasks: 1 of 4" in out + assert "1 error found" in out + + +def test_cli_analyze_slowest(run, config): + """`analyze --slowest N` prints exactly N rows.""" + record_session(config) + code, out = run("analyze", "--slowest", "2") + assert code == cli.OK + lines = [line for line in out.splitlines() if line.strip().endswith(("depth=0", "depth=1", "depth=2"))] + assert len(lines) == 2 + + +def test_cli_analyze_links(run, config): + """`analyze --links` prints the link count.""" + record_session(config) + code, out = run("analyze", "--links") + assert code == cli.OK + assert "cross-framework links: 0" in out + + +def test_cli_analyze_reports_a_miss(run, config): + """An unknown session prefix fails with a clear message.""" + record_session(config) + code, out = run("analyze", "definitely-not-a-session") + assert code == cli.FAILED + assert "No session matching" in out + + +# -- CLI ``analyze --compare`` ---------------------------------------------------- + +SECOND_SESSION = "sess-def" + + +def record_second_session(config): + """Capture a smaller error-free session under a second session id.""" + sid = SECOND_SESSION + fire(config, "SessionStart", source="startup", model="claude-opus-5", session_id=sid) + fire(config, "UserPromptSubmit", prompt="run it again", prompt_id="p1", session_id=sid) + fire(config, "PreToolUse", tool_name="Bash", tool_use_id="t1", tool_input={"command": "pytest -q"}, session_id=sid) + fire(config, "PostToolUse", tool_name="Bash", tool_use_id="t1", tool_response={"exit_code": 0}, session_id=sid) + fire(config, "Stop", last_assistant_message="All green.", session_id=sid) + fire(config, "SessionEnd", reason="clear", session_id=sid) + + +def test_cli_analyze_compare(run, config): + """`analyze --compare A B` prints per-activity count, duration, and error deltas.""" + record_session(config) + record_second_session(config) + session_a = ids.workflow_id_for("claude_code", SESSION) + session_b = ids.workflow_id_for("claude_code", SECOND_SESSION) + + code, out = run("analyze", "--compare", session_a[:8], session_b[:8]) + assert code == cli.OK + assert f"comparing: A={session_a} B={session_b}" in out + assert "tasks: 4 -> 2 (-2)" in out + assert "count 2 -> 1 (-1)" in out # Bash ran twice in A, once in B + assert "errors 50% -> 0%" in out + assert "only in A: Grep" in out + + +def test_cli_analyze_compare_reports_a_miss(run, config): + """An unknown session in either slot fails with a clear message.""" + record_session(config) + session_a = ids.workflow_id_for("claude_code", SESSION) + code, out = run("analyze", "--compare", session_a[:8], "definitely-not-a-session") + assert code == cli.FAILED + assert "No session matching" in out + + +def test_cli_analyze_compare_is_exclusive(run, config): + """--compare rejects a positional session and the single-session flags.""" + record_session(config) + record_second_session(config) + session_a = ids.workflow_id_for("claude_code", SESSION) + session_b = ids.workflow_id_for("claude_code", SECOND_SESSION) + + # argparse itself rejects a positional session next to --compare. + with pytest.raises(SystemExit): + cli.main(["analyze", session_a[:8], "--compare", session_a[:8], session_b[:8]]) + + code, out = run("analyze", "--compare", session_a[:8], session_b[:8], "--errors") + assert code == cli.FAILED + assert "--compare cannot be combined" in out diff --git a/tests/harness/test_claude_code.py b/tests/harness/test_claude_code.py new file mode 100644 index 00000000..b224c9d3 --- /dev/null +++ b/tests/harness/test_claude_code.py @@ -0,0 +1,269 @@ +"""End-to-end tests for the Claude Code adapter. + +These drive the adapter the way Claude Code does: one payload at a time, +each through :func:`~flowcept.agents.claude_code.claude_code_plugin.handle`, mimicking +the one-process-per-event model. +""" + +from __future__ import annotations + +import pytest + +from flowcept.agents.claude_code import claude_code_plugin as claude_code +from flowcept.agents.harness.vocab import ( + AGENT_SESSION, + AGENT_TOOL, + AI_MODEL_INVOCATION, + STATUS_ERROR, + STATUS_FINISHED, + STATUS_RUNNING, + SUBAGENT_SESSION, +) + +SESSION = "sess-abc" + + +def fire(config, event: str, **fields): + """Deliver one hook payload, as Claude Code would.""" + payload = {"hook_event_name": event, "session_id": SESSION, "cwd": "/tmp/proj", **fields} + return claude_code.handle(payload, config) + + +def by_type(records, kind): + return [r for r in records if r.get("type") == kind] + + +def one(records, **match): + found = [r for r in records if all(r.get(k) == v for k, v in match.items())] + assert len(found) == 1, f"expected exactly one record matching {match}, got {len(found)}" + return found[0] + + +@pytest.fixture +def session(config): + """A started session, ready for turns.""" + fire(config, "SessionStart", source="startup", model="claude-opus-5") + return config + + +# -- session lifecycle ------------------------------------------------------- + + +def test_session_start_emits_agent_and_workflow(config, buffer_records): + fire(config, "SessionStart", source="startup", model="claude-opus-5") + records = buffer_records() + + agent = one(records, type="agent") + assert agent["name"] == "claude_code:claude-opus-5" + + workflow = one(records, type="workflow") + assert workflow["subtype"] == AGENT_SESSION + assert workflow["status"] == STATUS_RUNNING + assert workflow["agent_id"] == agent["agent_id"] + + +def test_session_opens_once_across_processes(session, buffer_records): + """A second SessionStart (resume, or a racing hook) must not re-open.""" + fire(session, "SessionStart", source="resume", model="claude-opus-5") + assert len(by_type(buffer_records(), "agent")) == 1 + + +def test_session_without_id_is_dropped(config, buffer_records): + assert claude_code.handle({"hook_event_name": "SessionStart"}, config) == [] + assert buffer_records() == [] + + +def test_unknown_event_is_ignored(config, buffer_records): + assert claude_code.handle({"hook_event_name": "Nonesuch", "session_id": SESSION}, config) == [] + assert buffer_records() == [] + + +def test_session_end_supersedes_the_open_record(session, buffer_records): + """One workflow record per workflow, or Flowcept counts it as two runs.""" + fire(session, "SessionEnd", reason="clear") + + workflows = by_type(buffer_records(), "workflow") + assert len(workflows) == 1 + assert workflows[0]["status"] == STATUS_FINISHED + # The closing record has to restate what the dropped one said. + assert workflows[0]["used"]["session_id"] == SESSION + assert workflows[0]["used"]["start_reason"] == "startup" + assert workflows[0]["used"]["model"] == "claude-opus-5" + + +def test_session_records_the_harness_cwd_not_the_hooks(session, buffer_records): + """The hook process runs wherever it likes; only the harness cwd is real.""" + fire(session, "SessionEnd", reason="clear") + workflow = one(buffer_records(), type="workflow") + assert workflow["used"]["cwd"] == "/tmp/proj" + + +# -- turns and tools --------------------------------------------------------- + + +def test_turn_becomes_an_ai_model_invocation(session, buffer_records): + fire(session, "UserPromptSubmit", prompt="investigate the flaky test", prompt_id="p1") + fire(session, "Stop", last_assistant_message="It is a timing race.") + + turn = one(buffer_records(), subtype=AI_MODEL_INVOCATION) + assert turn["used"]["prompt"] == "investigate the flaky test" + assert turn["generated"]["response"] == "It is a timing race." + assert turn["status"] == STATUS_FINISHED + + +def test_tool_call_is_a_child_of_the_turn(session, buffer_records): + fire(session, "UserPromptSubmit", prompt="run the tests", prompt_id="p1") + fire(session, "PreToolUse", tool_name="Bash", tool_use_id="t1", tool_input={"command": "pytest"}) + fire(session, "PostToolUse", tool_name="Bash", tool_use_id="t1", tool_response={"exit_code": 0}) + fire(session, "Stop", last_assistant_message="green") + + records = buffer_records() + tool = one(records, subtype=AGENT_TOOL) + turn = one(records, subtype=AI_MODEL_INVOCATION) + + assert tool["activity_id"] == "Bash" + assert tool["used"]["command"] == "pytest" + assert tool["status"] == STATUS_FINISHED + # Pre and Post run in separate processes and must agree on the task id. + assert tool["parent_task_id"] == turn["task_id"] + + +def test_tool_emits_one_record_not_one_per_hook(session, buffer_records): + fire(session, "PreToolUse", tool_name="Read", tool_use_id="t1", tool_input={"file_path": "/x"}) + fire(session, "PostToolUse", tool_name="Read", tool_use_id="t1", tool_response={"ok": True}) + assert len(by_type(buffer_records(), "task")) == 1 + + +def test_failed_tool_is_recorded_as_error(session, buffer_records): + fire(session, "PreToolUse", tool_name="Bash", tool_use_id="t1", tool_input={"command": "false"}) + fire(session, "PostToolUseFailure", tool_name="Bash", tool_use_id="t1", error="exit status 1") + + tool = one(buffer_records(), subtype=AGENT_TOOL) + assert tool["status"] == STATUS_ERROR + assert tool["stderr"] == "exit status 1" + + +def test_post_without_pre_still_records(session, buffer_records): + """The Pre hook can be lost -- a crash, or capture enabled mid-session.""" + fire(session, "PostToolUse", tool_name="Glob", tool_use_id="t9", tool_response={"count": 2}) + assert one(buffer_records(), subtype=AGENT_TOOL)["activity_id"] == "Glob" + + +def test_new_prompt_closes_the_previous_turn(session, buffer_records): + """No Stop arrives when the user interrupts and types again.""" + fire(session, "UserPromptSubmit", prompt="first", prompt_id="p1") + fire(session, "UserPromptSubmit", prompt="second", prompt_id="p2") + + turns = [r for r in buffer_records() if r.get("subtype") == AI_MODEL_INVOCATION] + assert len(turns) == 1 + assert turns[0]["used"]["prompt"] == "first" + assert turns[0]["custom_metadata"]["close_reason"] == "superseded" + + +def test_open_turn_is_closed_at_session_end(session, buffer_records): + fire(session, "UserPromptSubmit", prompt="hello", prompt_id="p1") + fire(session, "SessionEnd", reason="clear") + + turn = one(buffer_records(), subtype=AI_MODEL_INVOCATION) + assert turn["custom_metadata"]["close_reason"] == "session_ended" + + +# -- subagents --------------------------------------------------------------- + + +def test_subagent_becomes_a_nested_workflow(session, buffer_records): + fire(session, "UserPromptSubmit", prompt="find it", prompt_id="p1") + fire(session, "PreToolUse", tool_name="Task", tool_use_id="t1", tool_input={"subagent_type": "Explore"}) + fire(session, "SubagentStart", agent_id="a1", agent_type="Explore", task="find the test") + fire(session, "PostToolUse", tool_name="Task", tool_use_id="t1", tool_response={"result": "found"}) + fire(session, "SubagentStop", agent_id="a1", agent_type="Explore", last_assistant_message="tests/test_net.py") + fire(session, "Stop", last_assistant_message="done") + + records = buffer_records() + parent = one(records, subtype=AGENT_SESSION) + sub = one(records, subtype=SUBAGENT_SESSION) + + assert sub["parent_workflow_id"] == parent["workflow_id"] + assert sub["status"] == STATUS_FINISHED + assert sub["used"]["agent_type"] == "Explore" + # Carried over from the superseded open record. + assert sub["used"]["prompt"] == "find the test" + assert sub["generated"]["response"] == "tests/test_net.py" + + +def test_subagent_tool_calls_land_in_the_subagent_workflow(session, buffer_records): + fire(session, "UserPromptSubmit", prompt="find it", prompt_id="p1") + fire(session, "PreToolUse", tool_name="Task", tool_use_id="t1", tool_input={"subagent_type": "Explore"}) + fire(session, "SubagentStart", agent_id="a1", agent_type="Explore") + fire(session, "PreToolUse", tool_name="Grep", tool_use_id="t2", tool_input={"pattern": "flaky"}, agent_id="a1") + fire(session, "PostToolUse", tool_name="Grep", tool_use_id="t2", tool_response={"matches": 3}, agent_id="a1") + fire(session, "SubagentStop", agent_id="a1", agent_type="Explore") + fire(session, "PostToolUse", tool_name="Task", tool_use_id="t1", tool_response={"result": "found"}) + + records = buffer_records() + sub = one(records, subtype=SUBAGENT_SESSION) + grep = one(records, activity_id="Grep") + task_tool = one(records, activity_id="Task") + + assert grep["workflow_id"] == sub["workflow_id"] + # The Task call itself belongs to the parent -- it is what spawned the sub. + assert task_tool["workflow_id"] == sub["parent_workflow_id"] + + +def test_session_counters_reflect_the_work(session, buffer_records): + fire(session, "UserPromptSubmit", prompt="go", prompt_id="p1") + fire(session, "PostToolUse", tool_name="Read", tool_use_id="t1", tool_response={}) + fire(session, "PostToolUseFailure", tool_name="Bash", tool_use_id="t2", error="boom") + fire(session, "SubagentStart", agent_id="a1", agent_type="Explore") + fire(session, "SubagentStop", agent_id="a1", agent_type="Explore") + fire(session, "SessionEnd", reason="clear") + + generated = one(buffer_records(), subtype=AGENT_SESSION)["generated"] + assert generated == {"turns": 1, "tool_calls": 2, "tool_errors": 1, "subagents": 1} + + +# -- capture safety ---------------------------------------------------------- + + +def test_secrets_are_redacted(session, buffer_records): + fire( + session, + "PostToolUse", + tool_name="Bash", + tool_use_id="t1", + tool_input={"command": 'curl -H "auth: sk-ant-abcdefghij0123456789"', "api_key": "hunter2"}, + tool_response={}, + ) + blob = str(one(buffer_records(), subtype=AGENT_TOOL)) + assert "sk-ant-abcdefghij0123456789" not in blob + assert "hunter2" not in blob + assert "«redacted»" in blob + + +def test_file_bodies_are_summarized_not_stored(session, buffer_records): + body = "\n".join(f"line {i}" for i in range(500)) + fire( + session, + "PostToolUse", + tool_name="Write", + tool_use_id="t1", + tool_input={"file_path": "/tmp/x.py", "content": body}, + tool_response={}, + ) + content = one(buffer_records(), subtype=AGENT_TOOL)["used"]["content"] + assert content["_summary"] is True + assert content["lines"] == 500 + assert len(str(content)) < len(body) + + +def test_capture_never_raises_into_the_harness(config): + """A malformed payload is a bug in the harness, not a reason to crash it.""" + assert claude_code.handle({"hook_event_name": "PreToolUse", "session_id": SESSION, "tool_input": object()}, config) + + +def test_hook_writes_nothing_to_stdout(session, capsys): + """Stdout on UserPromptSubmit is injected into the model's context.""" + fire(session, "UserPromptSubmit", prompt="hello", prompt_id="p1") + fire(session, "SessionEnd", reason="clear") + captured = capsys.readouterr() + assert captured.out == "" diff --git a/tests/harness/test_cli.py b/tests/harness/test_cli.py new file mode 100644 index 00000000..eab01dde --- /dev/null +++ b/tests/harness/test_cli.py @@ -0,0 +1,164 @@ +"""Tests for the ``flowcept-harness`` command.""" + +from __future__ import annotations + +import json + +import pytest + +from flowcept.agents.harness import cli + +from .test_claude_code import fire + + +@pytest.fixture +def run(config, monkeypatch, capsys): + """Invoke the CLI against the test's capture home.""" + monkeypatch.setenv("FLOWCEPT_HARNESS_HOME", str(config.home)) + + def _run(*argv: str): + code = cli.main(list(argv)) + return code, capsys.readouterr().out + + return _run + + +@pytest.fixture +def recorded(config): + fire(config, "SessionStart", source="startup", model="claude-opus-5") + fire(config, "UserPromptSubmit", prompt="fix it", prompt_id="p1") + fire(config, "PreToolUse", tool_name="Edit", tool_use_id="t1", tool_input={"file_path": "a.py"}) + fire(config, "PostToolUse", tool_name="Edit", tool_use_id="t1", tool_response={"ok": True}) + fire(config, "SessionEnd", reason="clear") + return config + + +def test_status_reports_the_capture_home(run, config): + code, out = run("status") + assert code == cli.OK + assert str(config.home) in out + + +def test_sessions_lists_nothing_before_capture(run): + code, out = run("sessions") + assert code == cli.OK + assert "No sessions" in out + + +def test_sessions_lists_a_captured_session(run, recorded): + code, out = run("sessions") + assert code == cli.OK + assert "FINISHED" in out + assert "claude_code session" in out + assert "tool_calls=1" in out + + +def test_show_lists_activity(run, recorded): + code, out = run("show") + assert code == cli.OK + assert "Edit" in out + assert "agent_turn" in out + + +def test_show_says_so_when_there_is_no_activity(run, config): + fire(config, "SessionStart", source="startup") + code, out = run("show") + assert code == cli.OK + assert "no activity" in out + + +def test_hook_subcommand_records(run, config, buffer_records, monkeypatch): + payload = {"session_id": "cli-1", "cwd": "/w", "source": "startup", "model": "m"} + monkeypatch.setattr("sys.stdin", _Stdin(json.dumps(payload))) + + code, _ = run("hook", "--event", "SessionStart") + assert code == cli.OK + assert any(r.get("type") == "workflow" for r in buffer_records()) + + +def test_hook_subcommand_routes_to_the_generic_adapter(run, buffer_records, monkeypatch): + payload = {"event": "session-start", "session_id": "cx-1", "cwd": "/repo"} + monkeypatch.setattr("sys.stdin", _Stdin(json.dumps(payload))) + + code, _ = run("hook", "--harness", "codex") + assert code == cli.OK + workflow = next(r for r in buffer_records() if r.get("type") == "workflow") + assert workflow["custom_metadata"]["harness"] == "codex" + + +def test_hook_never_fails_on_garbage(run, monkeypatch): + """A hook that exits non-zero is surfaced to the user mid-session.""" + monkeypatch.setattr("sys.stdin", _Stdin("this is not json")) + code, _ = run("hook", "--event", "SessionStart") + assert code == cli.OK + + +def test_repair_closes_a_dangling_session(run, config, buffer_records): + fire(config, "SessionStart", source="startup") + fire(config, "UserPromptSubmit", prompt="hi", prompt_id="p1") + + code, out = run("repair") + assert code == cli.OK + assert "1 session(s) repaired" in out + + workflow = next(r for r in buffer_records() if r.get("type") == "workflow") + assert workflow["status"] == "FINISHED" + assert workflow["custom_metadata"]["end_reason"] == "repaired" + + +def test_repair_is_idempotent(run, recorded): + code, out = run("repair") + assert code == cli.OK + assert "0 session(s) repaired" in out + + +def test_flush_dry_run_needs_no_backend(run, recorded, monkeypatch): + pytest.importorskip("flowcept") + published = [] + + class FakeMQ: + def bulk_publish(self, records): + published.extend(records) + + def stop(self): + pass + + monkeypatch.setattr("flowcept.commons.daos.mq_dao.mq_dao_base.MQDao.build", staticmethod(FakeMQ)) + code, out = run("flush", "--dry-run") + assert code == cli.OK + assert "would publish" in out + assert published == [], "a dry run must not publish" + + +def test_flush_publishes_and_can_remove(run, recorded, config, monkeypatch): + pytest.importorskip("flowcept") + published = [] + + class FakeMQ: + def bulk_publish(self, records): + published.extend(records) + + def stop(self): + pass + + monkeypatch.setattr("flowcept.commons.daos.mq_dao.mq_dao_base.MQDao.build", staticmethod(FakeMQ)) + code, out = run("flush", "--all", "--remove") + assert code == cli.OK + assert published, "records should reach the backend" + assert list(config.buffers_dir.glob("*.jsonl")) == [] + + +def test_install_prints_wiring(run): + code, out = run("install") + assert code == cli.OK + assert "PreToolUse" in out + + +class _Stdin: + """Minimal stdin stand-in; the adapter only ever calls read().""" + + def __init__(self, data: str): + self._data = data + + def read(self) -> str: + return self._data diff --git a/tests/harness/test_cross_capture.py b/tests/harness/test_cross_capture.py new file mode 100644 index 00000000..e2656e1a --- /dev/null +++ b/tests/harness/test_cross_capture.py @@ -0,0 +1,360 @@ +"""Tests for how the AI-harness capture and the agentic-framework plugins interoperate. + +Three real join mechanisms exist between the two capture systems: + +* a shared ``campaign_id``: the harness recorder stamps it on every buffer + record (from ``Config.campaign_id`` / ``FLOWCEPT_HARNESS_CAMPAIGN_ID``), and + the framework plugins stamp the same field on every emitted message, so the + two record sets join on it downstream; +* framework -> harness linking: a harness-emitted ``task_id`` can be passed as + ``_source_agent_id`` in a LangGraph initial state, and the LangGraph plugin + stores it as ``custom_metadata.source_agent_id`` on its records; +* harness -> framework linking: a framework-emitted task/agent id reaches the + harness recorder either as the ``flowcept_source_agent_id`` hook-payload key + or via the ``FLOWCEPT_HARNESS_SOURCE_AGENT_ID`` environment variable (the + payload key wins), and lands as ``source_agent_id`` on every turn, tool, and + LLM-call task the harness emits. ``SessionTracer`` takes the same value as + its ``source_agent_id`` argument. + +Framework emission is captured in memory by replacing ``BaseInterceptor`` with +a fake (the pattern used by tests/agents/plugins), so no MQ, MongoDB, or +network access is needed. Harness emission goes to the per-test JSONL buffer. +""" + +from __future__ import annotations + +import json + +import pytest + +import flowcept.flowceptor.adapters.base_interceptor as base_interceptor_module +from flowcept.agents.harness import SessionTracer, ids, prov +from flowcept.agents.harness.config import load_config +from flowcept.agents.harness.emit import Emitter +from flowcept.agents.harness.vocab import AGENT_TOOL, AI_MODEL_INVOCATION +from flowcept.agents.langgraph.langgraph_plugin import FlowceptLangGraphPlugin + +from .test_claude_code import fire + +pytest.importorskip("langgraph") +pytest.importorskip("langchain_core") + + +class _CapturingInterceptor: + """In-memory stand-in for BaseInterceptor that records all emissions.""" + + instances: list = [] + + def __init__(self, plugin_key=None, kind=None): + """Record construction and initialize empty capture buffers.""" + self.kind = kind + self.telemetry_capture = None + self.stopped = False + self.task_messages: list[dict] = [] + self.workflow_messages: list = [] + type(self).instances.append(self) + + def start(self, bundle_exec_id, check_safe_stops=True): + """Mark the interceptor as started.""" + return self + + def stop(self, check_safe_stops=True): + """Mark the interceptor as stopped.""" + self.stopped = True + + def send_workflow_message(self, workflow_obj): + """Capture a WorkflowObject emission.""" + self.workflow_messages.append(workflow_obj) + + def intercept(self, task_dict): + """Capture a task message emission.""" + self.task_messages.append(task_dict) + + +@pytest.fixture() +def captured(monkeypatch): + """Patch BaseInterceptor with the capturing fake and return its instance list.""" + _CapturingInterceptor.instances = [] + monkeypatch.setattr(base_interceptor_module, "BaseInterceptor", _CapturingInterceptor) + return _CapturingInterceptor.instances + + +def _build_graph(): + """Build a small local StateGraph of plain python-function nodes.""" + from typing import TypedDict + + from langgraph.graph import END, START, StateGraph + + class _State(TypedDict, total=False): + value: int + _source_agent_id: str + + def _add_one(state): + return {"value": state["value"] + 1} + + def _double(state): + return {"value": state["value"] * 2} + + builder = StateGraph(_State) + builder.add_node("add_one", _add_one) + builder.add_node("double", _double) + builder.add_edge(START, "add_one") + builder.add_edge("add_one", "double") + builder.add_edge("double", END) + return builder.compile() + + +def _run_graph(captured, initial_state, campaign_id=None): + """Invoke the local graph under a fresh LangGraph plugin; return its fake.""" + plugin_config = {"workflow_name": "cross-wf", "performance_tracking": False} + if campaign_id: + plugin_config["campaign_id"] = campaign_id + plugin = FlowceptLangGraphPlugin(config=plugin_config) + plugin.start() + fake = captured[-1] + try: + graph = _build_graph() + graph.invoke(initial_state, config={"callbacks": [plugin.callback_handler]}) + finally: + plugin.stop() + return fake + + +def _tasks(records, subtype=None): + """Return the task records, optionally filtered by subtype.""" + return [r for r in records if r.get("type") == "task" and (subtype is None or r.get("subtype") == subtype)] + + +# -- source_agent_id on harness records ---------------------------------------- + + +def test_task_record_carries_source_agent_id_into_the_buffer(config, buffer_records): + """A harness task record built with a source agent id keeps it end-to-end. + + This exercises the record-builder/emitter level directly; the recorder, + hook adapters, and tracer paths are covered by the tests below. + """ + workflow_id = ids.workflow_id_for("sdk_agent", "link-run") + record = prov.task_record( + task_id=ids.tool_task_id(workflow_id, "t1"), + workflow_id=workflow_id, + activity_id="dispatch", + subtype=AGENT_TOOL, + source_agent_id="framework-agent-9", + ) + Emitter(config, workflow_id).emit(record) + + (buffered,) = buffer_records() + assert buffered["source_agent_id"] == "framework-agent-9" + assert buffered["task_id"] == record["task_id"] + + +def test_recorder_emitted_tasks_have_no_source_agent_id_when_unset(config, buffer_records): + """With no payload key, env var, or tracer argument, the field stays absent. + + None-valued keys are dropped by ``prov._clean``, so a harness run that was + not given a source agent id never emits the field at all. + """ + with SessionTracer("sdk_agent", "plain-run", config=config) as tracer: + tracer.prompt("hello") + call = tracer.tool_start("search", {"q": "x"}) + tracer.tool_end(call, tool_response={"hits": 1}) + tracer.turn_end("done") + + task_records = _tasks(buffer_records()) + assert task_records + for record in task_records: + assert "source_agent_id" not in record + + +def test_payload_key_sets_source_agent_id_on_harness_tasks(config, buffer_records): + """The ``flowcept_source_agent_id`` hook-payload key lands on turn and tool tasks.""" + fire(config, "SessionStart", source="startup", model="claude-opus-5") + fire(config, "UserPromptSubmit", prompt="go", prompt_id="p1") + fire(config, "PreToolUse", tool_name="Bash", tool_use_id="t1", tool_input={"command": "ls"}) + fire( + config, + "PostToolUse", + tool_name="Bash", + tool_use_id="t1", + tool_response={"ok": True}, + flowcept_source_agent_id="fw-task-42", + ) + fire(config, "Stop", last_assistant_message="done", flowcept_source_agent_id="fw-task-42") + fire(config, "SessionEnd", reason="clear") + + records = buffer_records() + (tool,) = _tasks(records, AGENT_TOOL) + (turn,) = _tasks(records, AI_MODEL_INVOCATION) + assert tool["source_agent_id"] == "fw-task-42" + assert turn["source_agent_id"] == "fw-task-42" + + +def test_env_source_agent_id_reaches_harness_tasks(monkeypatch, tmp_path): + """FLOWCEPT_HARNESS_SOURCE_AGENT_ID flows through load_config onto every task.""" + monkeypatch.setenv("FLOWCEPT_HARNESS_SOURCE_AGENT_ID", "fw-env-7") + monkeypatch.setenv("FLOWCEPT_HARNESS_HOME", str(tmp_path / "env-home")) + env_config = load_config() + assert env_config.source_agent_id == "fw-env-7" + + with SessionTracer("sdk_agent", "env-src-run", config=env_config) as tracer: + tracer.prompt("hi") + call = tracer.tool_start("search", {"q": "x"}) + tracer.tool_end(call, tool_response={"hits": 1}) + tracer.turn_end("done") + + records = [] + for path in sorted(env_config.buffers_dir.glob("*.jsonl")): + records.extend(json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()) + task_records = _tasks(records) + assert task_records + for record in task_records: + assert record["source_agent_id"] == "fw-env-7" + + +def test_payload_key_wins_over_env_source_agent_id(config, buffer_records): + """When both are set, the payload key beats the env-configured value.""" + config.source_agent_id = "fw-from-env" # what load_config would have set + fire(config, "PreToolUse", tool_name="Bash", tool_use_id="t1", tool_input={"command": "ls"}) + fire( + config, + "PostToolUse", + tool_name="Bash", + tool_use_id="t1", + tool_response={"ok": True}, + flowcept_source_agent_id="fw-from-payload", + ) + fire(config, "Stop", last_assistant_message="done") + + records = buffer_records() + (tool,) = _tasks(records, AGENT_TOOL) + assert tool["source_agent_id"] == "fw-from-payload" + # An event without the payload key still falls back to the env value. + (turn,) = _tasks(records, AI_MODEL_INVOCATION) + assert turn["source_agent_id"] == "fw-from-env" + + +# -- campaign_id as the cross-system join key ----------------------------------- + + +def test_shared_campaign_id_joins_harness_and_langgraph_records(config, buffer_records, captured): + """The same campaign_id on both captures lands on every record of each.""" + config.campaign_id = "camp-joint" + with SessionTracer("sdk_agent", "camp-run", config=config) as tracer: + tracer.prompt("plan") + tracer.turn_end("planned") + + fake = _run_graph(captured, {"value": 1}, campaign_id="camp-joint") + + harness_records = buffer_records() + assert harness_records + for record in harness_records: + if record.get("type") in ("task", "workflow", "agent"): + assert record["campaign_id"] == "camp-joint" + + assert fake.workflow_messages[0].campaign_id == "camp-joint" + assert fake.task_messages + for task in fake.task_messages: + assert task["campaign_id"] == "camp-joint" + + +def test_env_campaign_id_reaches_harness_records(monkeypatch, tmp_path): + """FLOWCEPT_HARNESS_CAMPAIGN_ID flows through load_config onto the buffer.""" + monkeypatch.setenv("FLOWCEPT_HARNESS_CAMPAIGN_ID", "camp-env") + monkeypatch.setenv("FLOWCEPT_HARNESS_HOME", str(tmp_path / "env-home")) + env_config = load_config() + assert env_config.campaign_id == "camp-env" + + with SessionTracer("sdk_agent", "env-run", config=env_config) as tracer: + tracer.prompt("hi") + tracer.turn_end("done") + + records = [] + for path in sorted(env_config.buffers_dir.glob("*.jsonl")): + records.extend(json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()) + assert records + assert {r["campaign_id"] for r in records if r.get("type") in ("task", "workflow")} == {"camp-env"} + + +# -- coexistence ----------------------------------------------------------------- + + +def test_harness_and_langgraph_captures_coexist_in_one_process(config, buffer_records, captured): + """Both captures running at once each produce intact, non-colliding records.""" + plugin = FlowceptLangGraphPlugin(config={"workflow_name": "co-wf", "performance_tracking": False}) + plugin.start() + fake = captured[-1] + try: + with SessionTracer("sdk_agent", "co-run", config=config) as tracer: + tracer.prompt("run the graph") + graph = _build_graph() + result = graph.invoke({"value": 3}, config={"callbacks": [plugin.callback_handler]}) + call = tracer.tool_start("invoke_graph", {"value": 3}) + tracer.tool_end(call, tool_response={"value": result["value"]}) + tracer.turn_end(f"graph returned {result['value']}") + finally: + plugin.stop() + + # The harness side is complete: buffer parses, and counts are as expected. + harness_records = buffer_records() + assert len(_tasks(harness_records, AI_MODEL_INVOCATION)) == 1 + (tool,) = _tasks(harness_records, AGENT_TOOL) + assert tool["generated"]["value"] == 8 + session = [r for r in harness_records if r.get("type") == "workflow"][-1] + assert session["status"] == "FINISHED" + assert session["generated"] == {"turns": 1, "tool_calls": 1} + + # The framework side is complete too, and untouched by the harness capture. + graph_tasks = [t for t in fake.task_messages if t["subtype"] == "langgraph_graph"] + node_names = {t["activity_id"] for t in fake.task_messages if t["subtype"] == "langgraph_node"} + assert len(graph_tasks) == 1 + assert {"add_one", "double"}.issubset(node_names) + + # No identifier from one system leaks into or collides with the other. + harness_ids = {r.get("task_id") for r in _tasks(harness_records)} + harness_ids |= {r["workflow_id"] for r in harness_records if r.get("type") == "workflow"} + framework_ids = {t["task_id"] for t in fake.task_messages} + framework_ids |= {w.workflow_id for w in fake.workflow_messages} + assert not harness_ids & framework_ids + + +# -- framework -> harness linking ------------------------------------------------- + + +def test_harness_tool_task_id_round_trips_through_a_langgraph_run(config, buffer_records, captured): + """A real harness-emitted task_id passed as _source_agent_id lands on all records.""" + with SessionTracer("sdk_agent", "link-src", config=config) as tracer: + call = tracer.tool_start("prepare_input", {"n": 1}) + tracer.tool_end(call, tool_response={"ready": True}) + + (harness_tool,) = _tasks(buffer_records(), AGENT_TOOL) + source_task_id = harness_tool["task_id"] + + fake = _run_graph(captured, {"value": 1, "_source_agent_id": source_task_id}) + + graph_task = next(t for t in fake.task_messages if t["subtype"] == "langgraph_graph") + assert graph_task["custom_metadata"]["source_agent_id"] == source_task_id + for node_task in (t for t in fake.task_messages if t["subtype"] == "langgraph_node"): + assert node_task["custom_metadata"]["source_agent_id"] == source_task_id + + +# -- harness -> framework linking ------------------------------------------------- + + +def test_framework_task_id_round_trips_into_harness_records(config, buffer_records, captured): + """A real framework-emitted task_id lands as source_agent_id on harness tasks.""" + fake = _run_graph(captured, {"value": 1}) + graph_task = next(t for t in fake.task_messages if t["subtype"] == "langgraph_graph") + framework_task_id = graph_task["task_id"] + + with SessionTracer("sdk_agent", "link-back", config=config, source_agent_id=framework_task_id) as tracer: + tracer.prompt("analyze the graph run") + call = tracer.tool_start("inspect", {"target": "graph"}) + tracer.tool_end(call, tool_response={"ok": True}) + tracer.turn_end("done") + + records = buffer_records() + (tool,) = _tasks(records, AGENT_TOOL) + (turn,) = _tasks(records, AI_MODEL_INVOCATION) + assert tool["source_agent_id"] == framework_task_id + assert turn["source_agent_id"] == framework_task_id diff --git a/tests/harness/test_flowcept_interop.py b/tests/harness/test_flowcept_interop.py new file mode 100644 index 00000000..dad61f14 --- /dev/null +++ b/tests/harness/test_flowcept_interop.py @@ -0,0 +1,100 @@ +"""The buffer must be readable by Flowcept with no conversion step. + +This is the whole premise of the JSONL format choice, so it is tested against +the real library rather than a stub. Skipped when flowcept is not installed -- +the capture path does not depend on it. +""" + +from __future__ import annotations + +import pytest + +from flowcept.agents.harness.vocab import AGENT_SESSION, SUBAGENT_SESSION + +from .test_claude_code import fire + +flowcept = pytest.importorskip("flowcept") + + +@pytest.fixture +def captured_session(config): + """A realistic session: a turn, tools, a subagent, and a failure.""" + fire(config, "SessionStart", source="startup", model="claude-opus-5") + fire(config, "UserPromptSubmit", prompt="investigate the flaky test", prompt_id="p1") + fire(config, "PreToolUse", tool_name="Task", tool_use_id="t1", tool_input={"subagent_type": "Explore"}) + fire(config, "SubagentStart", agent_id="a1", agent_type="Explore", task="find the test") + fire(config, "PreToolUse", tool_name="Grep", tool_use_id="t2", tool_input={"pattern": "flaky"}, agent_id="a1") + fire(config, "PostToolUse", tool_name="Grep", tool_use_id="t2", tool_response={"matches": 3}, agent_id="a1") + fire(config, "SubagentStop", agent_id="a1", agent_type="Explore", last_assistant_message="tests/test_net.py") + fire(config, "PostToolUse", tool_name="Task", tool_use_id="t1", tool_response={"result": "found"}) + fire(config, "PreToolUse", tool_name="Bash", tool_use_id="t3", tool_input={"command": "pytest"}) + fire(config, "PostToolUseFailure", tool_name="Bash", tool_use_id="t3", error="exit status 1") + fire(config, "Stop", last_assistant_message="The flake is a timing race.") + fire(config, "SessionEnd", reason="clear") + return next(iter(config.buffers_dir.glob("*.jsonl"))) + + +def test_flowcept_loads_the_buffer_without_conversion(captured_session): + from flowcept.report.loaders import read_jsonl, split_records + + records, skipped = read_jsonl(captured_session) + assert skipped == 0, "every line must be a record Flowcept understands" + + dataset = split_records(records) + workflows = dataset["workflows"] + subtypes = sorted(w.get("subtype") for w in workflows) + + # One record per workflow -- the session and its one subagent. Two records + # for either would be counted as two separate runs. + assert subtypes == [AGENT_SESSION, SUBAGENT_SESSION] + + +def test_generate_report_sees_the_right_shape(captured_session, tmp_path): + from flowcept import Flowcept + + out = tmp_path / "card.md" + stats = Flowcept.generate_report( + report_type="workflow_card", + input_jsonl_path=str(captured_session), + format="markdown", + output_path=str(out), + ) + + assert stats["skipped_lines"] == 0 + assert stats["n_workflows"] == 2 # session + subagent + assert stats["n_tasks"] == 4 # turn + Task + Grep + Bash + + card = out.read_text(encoding="utf-8") + assert "subagent:Explore" in card + assert card.count("**Workflow ID:**") == 2 + + +def test_an_sdk_captured_session_reports_the_same_way(config, tmp_path): + """The in-process path must produce the same shape as the hook path.""" + from flowcept import Flowcept + + from flowcept.agents.harness import SessionTracer + + with SessionTracer("claude_agent_sdk", "s1", config=config, model="claude-opus-5") as tracer: + tracer.prompt("fix the tests") + ref = tracer.subagent_start("Explore", prompt="find them") + call = tracer.tool_start("Grep", {"pattern": "def test"}, agent_ref=ref) + tracer.tool_end(call, tool_response={"matches": 12}, agent_ref=ref) + tracer.subagent_stop(ref, response="found 12") + edit = tracer.tool_start("Edit", {"file_path": "a.py"}) + tracer.tool_end(edit, tool_response="ok") + tracer.turn_end("done", usage={"input_tokens": 500}) + + buffer = next(iter(config.buffers_dir.glob("*.jsonl"))) + out = tmp_path / "sdk_card.md" + stats = Flowcept.generate_report( + report_type="workflow_card", + input_jsonl_path=str(buffer), + format="markdown", + output_path=str(out), + ) + + assert stats["skipped_lines"] == 0 + assert stats["n_workflows"] == 2 # session + subagent + assert stats["n_tasks"] == 3 # turn + Grep + Edit + assert "subagent:Explore" in out.read_text(encoding="utf-8") diff --git a/tests/harness/test_generic.py b/tests/harness/test_generic.py new file mode 100644 index 00000000..9bc6e2f3 --- /dev/null +++ b/tests/harness/test_generic.py @@ -0,0 +1,253 @@ +"""Tests for the profile-driven adapter used by non-Claude-Code harnesses.""" + +from __future__ import annotations + +import json + +import pytest + +from flowcept.agents.cli_harness import cli_harness_plugin as generic +from flowcept.agents.cli_harness.cli_harness_plugin import PROFILE_DIR, Profile, _dig, to_event +from flowcept.agents.harness.vocab import AGENT_TOOL, AI_MODEL_INVOCATION, EventKind + + +def test_dig_walks_dicts_and_lists(): + """_dig descends dotted paths through dicts and list indices.""" + payload = {"a": {"b": [{"c": 1}]}} + assert _dig(payload, "a.b.0.c") == 1 + assert _dig(payload, "a.b.9.c") is None + assert _dig(payload, "a.missing.c") is None + + +def test_event_names_match_regardless_of_spelling(): + """Event-name matching ignores case and separators.""" + profile = Profile("x") + for spelling in ("PreToolUse", "pre_tool_use", "pre-tool-use", "tool.before"): + assert profile.kind_for(spelling) == EventKind.TOOL_PRE, spelling + + +def test_unknown_harness_works_without_a_profile(config, buffer_records): + """Default field names should carry a harness nobody has written up yet.""" + generic.handle( + {"event": "session_start", "session_id": "s1", "cwd": "/w", "model": "gpt-5"}, + config, + harness="mystery", + ) + generic.handle( + { + "event": "tool_end", + "session_id": "s1", + "toolName": "shell", + "toolCallId": "c1", + "arguments": {"cmd": "ls"}, + "result": {"ok": True}, + }, + config, + harness="mystery", + ) + + tool = next(r for r in buffer_records() if r.get("subtype") == AGENT_TOOL) + assert tool["activity_id"] == "shell" + assert tool["used"]["cmd"] == "ls" + assert tool["custom_metadata"]["harness"] == "mystery" + + +def test_unmapped_event_is_ignored(config, buffer_records): + """An event name mapped to no kind produces no records.""" + assert generic.handle({"event": "heartbeat", "session_id": "s1"}, config) == [] + assert buffer_records() == [] + + +def test_missing_event_name_is_ignored(config): + """A payload with no recognizable event-name key produces no records.""" + assert generic.handle({"session_id": "s1"}, config) == [] + + +@pytest.mark.parametrize("name", ["codex", "gemini", "cursor", "opencode"]) +def test_shipped_profiles_are_valid(name): + """Every shipped profile maps its events onto kinds the recorder handles.""" + data = json.loads((PROFILE_DIR / f"{name}.json").read_text(encoding="utf-8")) + assert data["harness"] + assert data["events"], "a profile with no event map adds nothing" + + profile = Profile.load(name, name) + # Every declared event must resolve to a kind the recorder handles. + for source_name, kind in data["events"].items(): + assert kind in EventKind.ALL, f"{name}: {source_name} -> unknown kind {kind}" + assert profile.kind_for(source_name) == kind + + +def test_profile_falls_back_when_file_is_missing(): + """A missing profile file falls back to the built-in defaults.""" + profile = Profile.load("no-such-profile", "somewhere") + assert profile.harness == "somewhere" + assert profile.kind_for("SessionStart") == EventKind.SESSION_START + + +def test_codex_nested_fields(config, buffer_records): + """The codex profile digs tool detail out of the nested invocation object.""" + generic.handle({"event": "session-start", "session_id": "c1", "cwd": "/repo"}, config, harness="codex") + generic.handle( + { + "event": "mcp-tool-call-end", + "session_id": "c1", + "call_id": "call-1", + "invocation": {"tool": "fetch", "arguments": {"url": "https://x"}}, + "output": {"status": 200}, + }, + config, + harness="codex", + ) + + tool = next(r for r in buffer_records() if r.get("subtype") == AGENT_TOOL) + assert tool["activity_id"] == "fetch" + assert tool["used"]["url"] == "https://x" + assert tool["generated"]["status"] == 200 + + +def test_opencode_dotted_paths(config, buffer_records): + """The opencode profile resolves dotted paths like properties.info.id.""" + generic.handle( + {"type": "session.created", "properties": {"info": {"id": "o1"}}, "directory": "/repo"}, + config, + harness="opencode", + ) + generic.handle( + { + "type": "tool.execute.after", + "sessionID": "o1", + "tool": "edit", + "callID": "k1", + "args": {"file": "a.py"}, + "output": "done", + }, + config, + harness="opencode", + ) + + records = buffer_records() + assert next(r for r in records if r.get("subtype") == AGENT_TOOL)["activity_id"] == "edit" + + +def test_cursor_workspace_root_is_indexed(config, buffer_records): + """The cursor profile indexes the first entry of workspace_roots as cwd.""" + generic.handle( + { + "hook_event_name": "start", + "conversation_id": "x1", + "workspace_roots": ["/home/me/repo", "/other"], + }, + config, + harness="cursor", + ) + workflow = next(r for r in buffer_records() if r.get("type") == "workflow") + assert workflow["used"]["cwd"] == "/home/me/repo" + + +def test_gemini_camelcase_session_fields(config, buffer_records): + """The gemini profile reads camelCase keys and renames the harness to gemini_cli.""" + generic.handle( + { + "event": "SessionStart", + "sessionId": "g1", + "workspaceDir": "/repo", + "model": "gemini-2.5-pro", + }, + config, + harness="gemini", + ) + workflow = next(r for r in buffer_records() if r.get("type") == "workflow") + assert workflow["used"]["cwd"] == "/repo" + assert workflow["used"]["model"] == "gemini-2.5-pro" + assert workflow["custom_metadata"]["harness"] == "gemini_cli" + + +def test_gemini_nested_toolcall_object(config, buffer_records): + """Tool detail nested under gemini's toolCall object lands on one tool task.""" + generic.handle({"event": "SessionStart", "sessionId": "g1"}, config, harness="gemini") + generic.handle( + { + "event": "BeforeToolCall", + "sessionId": "g1", + "toolCall": {"name": "run_shell_command", "callId": "t1", "args": {"command": "ls"}}, + }, + config, + harness="gemini", + ) + generic.handle( + { + "event": "AfterToolCall", + "sessionId": "g1", + "toolCall": {"name": "run_shell_command", "callId": "t1", "response": {"exit_code": 0}}, + }, + config, + harness="gemini", + ) + + tools = [r for r in buffer_records() if r.get("subtype") == AGENT_TOOL] + assert len(tools) == 1 # the pre/post pair pairs up on toolCall.callId + assert tools[0]["activity_id"] == "run_shell_command" + assert tools[0]["used"]["command"] == "ls" + assert tools[0]["generated"]["exit_code"] == 0 + assert tools[0]["custom_metadata"]["duration_known"] is True + + +def test_gemini_prompt_and_model_response_form_a_turn(config, buffer_records): + """UserPromptSubmit opens a turn that ModelResponse closes with the answer.""" + generic.handle({"event": "SessionStart", "sessionId": "g1"}, config, harness="gemini") + generic.handle({"event": "UserPromptSubmit", "sessionId": "g1", "prompt": "refactor"}, config, harness="gemini") + generic.handle({"event": "ModelResponse", "sessionId": "g1", "responseText": "done"}, config, harness="gemini") + + turn = next(r for r in buffer_records() if r.get("subtype") == AI_MODEL_INVOCATION) + assert turn["used"]["prompt"] == "refactor" + assert turn["generated"]["response"] == "done" + assert turn["custom_metadata"]["harness"] == "gemini_cli" + + +def test_gemini_tool_error_is_a_failed_task(config, buffer_records): + """ToolCallError becomes an ERROR task with the nested toolCall.error as stderr.""" + generic.handle({"event": "SessionStart", "sessionId": "g1"}, config, harness="gemini") + generic.handle( + { + "event": "ToolCallError", + "sessionId": "g1", + "toolCall": {"name": "write_file", "callId": "t9", "error": "permission denied"}, + }, + config, + harness="gemini", + ) + + tool = next(r for r in buffer_records() if r.get("subtype") == AGENT_TOOL) + assert tool["activity_id"] == "write_file" + assert tool["status"] == "ERROR" + assert tool["stderr"] == "permission denied" + + +def test_prompt_and_tool_fields_do_not_collide(config): + """`input` means a prompt on a turn and arguments on a tool call.""" + profile = Profile("x") + prompt_event = to_event( + {"event": "prompt", "session_id": "s", "input": "do the thing"}, + harness="x", + profile=profile, + ) + tool_event = to_event( + {"event": "tool_start", "session_id": "s", "name": "grep", "input": {"pattern": "x"}}, + harness="x", + profile=profile, + ) + assert prompt_event.prompt == "do the thing" + assert prompt_event.tool_input is None + assert tool_event.tool_input == {"pattern": "x"} + assert tool_event.prompt is None + + +def test_generic_turn_is_recorded(config, buffer_records): + """A prompt/turn_end pair from an unprofiled harness becomes one turn task.""" + generic.handle({"event": "session_start", "session_id": "s1"}, config, harness="h") + generic.handle({"event": "prompt", "session_id": "s1", "prompt": "hi"}, config, harness="h") + generic.handle({"event": "turn_end", "session_id": "s1", "response": "hello"}, config, harness="h") + + turn = next(r for r in buffer_records() if r.get("subtype") == AI_MODEL_INVOCATION) + assert turn["used"]["prompt"] == "hi" + assert turn["generated"]["response"] == "hello" diff --git a/tests/harness/test_mcp_server.py b/tests/harness/test_mcp_server.py new file mode 100644 index 00000000..fec9ae48 --- /dev/null +++ b/tests/harness/test_mcp_server.py @@ -0,0 +1,146 @@ +"""Tests for the MCP provenance server. + +The tools are exercised through the server's own registry rather than by +calling the local functions, so a rename or signature change in the SDK's +decorator is caught here. +""" + +from __future__ import annotations + +import pytest + +from .test_claude_code import fire + +pytest.importorskip("mcp") + +from flowcept.agents.harness import mcp_server # noqa: E402 + + +@pytest.fixture +def server(config): + fire(config, "SessionStart", source="startup", model="claude-opus-5") + fire(config, "UserPromptSubmit", prompt="fix the flake", prompt_id="p1") + fire(config, "PreToolUse", tool_name="Bash", tool_use_id="t1", tool_input={"command": "pytest -q"}) + fire(config, "PostToolUse", tool_name="Bash", tool_use_id="t1", tool_response={"exit_code": 0}) + fire(config, "PreToolUse", tool_name="Bash", tool_use_id="t2", tool_input={"command": "ruff check"}) + fire(config, "PostToolUseFailure", tool_name="Bash", tool_use_id="t2", error="1 error found") + fire(config, "SubagentStart", agent_id="a1", agent_type="Explore") + fire(config, "PreToolUse", tool_name="Grep", tool_use_id="t3", tool_input={"pattern": "flaky"}, agent_id="a1") + fire(config, "PostToolUse", tool_name="Grep", tool_use_id="t3", tool_response={"matches": 2}, agent_id="a1") + fire(config, "SubagentStop", agent_id="a1", agent_type="Explore") + fire(config, "Stop", last_assistant_message="Fixed the race.") + fire(config, "SessionEnd", reason="clear") + return mcp_server.build_server(config) + + +def call(server, name: str, **kwargs): + """Invoke a registered tool by name, as a client would.""" + fn = getattr(server, "_tool_functions", {}).get(name) + if fn is None: + # Both SDK generations keep the undecorated callable reachable; fall + # back to the module-level lookup used by the tool manager. + manager = getattr(server, "_tool_manager", None) + tool = manager.get_tool(name) if manager else None + fn = getattr(tool, "fn", None) + assert fn is not None, f"tool {name!r} is not registered" + return fn(**kwargs) + + +def test_expected_tools_are_registered(server): + manager = getattr(server, "_tool_manager", None) + assert manager is not None, "SDK no longer exposes a tool manager" + names = {t.name for t in manager.list_tools()} + assert names == { + "list_sessions", + "get_session", + "search_tool_calls", + "session_stats", + "record_event", + "generate_report", + "analyze_session", + "analyze_errors", + "find_slowest", + "cross_links", + } + + +def test_every_tool_has_a_description(server): + for tool in server._tool_manager.list_tools(): + assert tool.description, f"{tool.name} has no description for the model to read" + + +def test_list_sessions(server): + sessions = call(server, "list_sessions") + assert len(sessions) == 1 + entry = sessions[0] + assert entry["harness"] == "claude_code" + assert entry["status"] == "FINISHED" + assert entry["totals"]["tool_calls"] == 3 + assert entry["totals"]["tool_errors"] == 1 + + +def test_get_session_separates_subagent_work(server): + result = call(server, "get_session") + assert result["status"] == "FINISHED" + assert [s["name"] for s in result["subagents"]] == ["subagent:Explore"] + + grep = next(a for a in result["activity"] if a["name"] == "Grep") + assert grep["in_subagent"] == "subagent:Explore" + bash = next(a for a in result["activity"] if a["name"] == "Bash") + assert "in_subagent" not in bash + + +def test_get_session_omits_io_by_default(server): + activity = call(server, "get_session")["activity"] + assert all("used" not in a for a in activity) + with_io = call(server, "get_session", include_io=True)["activity"] + assert any(a.get("used") for a in with_io) + + +def test_get_session_reports_a_miss(server): + assert "error" in call(server, "get_session", session="nope") + + +def test_search_by_status_finds_failures(server): + hits = call(server, "search_tool_calls", status="ERROR") + assert len(hits) == 1 + assert hits[0]["error"] == "1 error found" + + +def test_search_by_content(server): + hits = call(server, "search_tool_calls", contains="ruff") + assert len(hits) == 1 + assert hits[0]["tool"] == "Bash" + + +def test_search_respects_the_limit(server): + assert len(call(server, "search_tool_calls", limit=1)) == 1 + + +def test_session_stats_aggregates_by_tool(server): + stats = call(server, "session_stats") + assert stats["by_tool"]["Bash"] == {"calls": 2, "errors": 1, "seconds": pytest.approx(0, abs=1)} + assert stats["by_tool"]["Grep"]["calls"] == 1 + + +def test_record_event_writes_provenance(config): + server = mcp_server.build_server(config) + assert call(server, "record_event", kind="session_start", session_id="m1", harness="my_agent")["recorded"] + result = call( + server, + "record_event", + kind="tool_post", + session_id="m1", + harness="my_agent", + tool_name="query_db", + tool_input={"sql": "select 1"}, + ) + assert result["recorded"] == 1 + + hits = call(server, "search_tool_calls", tool_name="query_db") + assert hits[0]["used"]["sql"] == "select 1" + + +def test_record_event_rejects_an_unknown_kind(config): + server = mcp_server.build_server(config) + assert "error" in call(server, "record_event", kind="nonsense", session_id="m1") diff --git a/tests/harness/test_otel.py b/tests/harness/test_otel.py new file mode 100644 index 00000000..575983b4 --- /dev/null +++ b/tests/harness/test_otel.py @@ -0,0 +1,258 @@ +"""Tests for OpenTelemetry ingest.""" + +from __future__ import annotations + +import json + +import pytest + +from flowcept.agents.otel import otel_plugin as otel +from flowcept.agents.harness.vocab import AGENT_TOOL, AI_MODEL_INVOCATION, HARNESS_EVENT + +CONVERSATION = "gen_ai.conversation.id" + + +def span(**attributes): + """Build a minimal console-exporter-shaped span.""" + return { + "name": attributes.pop("_name", "span"), + "attributes": {CONVERSATION: "conv-1", **attributes}, + "start_time": 1_700_000_000_000_000_000, + "end_time": 1_700_000_001_500_000_000, + "context": {"span_id": "abc123"}, + "status": {"status_code": attributes.pop("_status", "OK")}, + } + + +def test_tool_span_becomes_a_tool_task(config, buffer_records): + """A tool-execution span maps to an agent_tool task with parsed arguments.""" + assert otel.ingest_spans( + [ + span( + **{ + "gen_ai.operation.name": "execute_tool", + "gen_ai.tool.name": "search_docs", + "gen_ai.tool.call.arguments": '{"query": "provenance"}', + "gen_ai.tool.call.result": '{"hits": 3}', + } + ) + ], + config, + ) + + tool = next(r for r in buffer_records() if r.get("subtype") == AGENT_TOOL) + assert tool["activity_id"] == "search_docs" + # JSON-encoded attributes are parsed back into real fields. + assert tool["used"]["query"] == "provenance" + assert tool["generated"]["hits"] == 3 + assert tool["ended_at"] - tool["started_at"] == pytest.approx(1.5, abs=0.01) + + +def test_model_span_becomes_a_model_invocation(config, buffer_records): + """A chat span maps to an ai_model_invocation task with model and usage.""" + otel.ingest_spans( + [ + span( + **{ + "gen_ai.operation.name": "chat", + "gen_ai.request.model": "claude-opus-5", + "gen_ai.usage.input_tokens": 1200, + "gen_ai.usage.output_tokens": 340, + } + ) + ], + config, + ) + + call = next(r for r in buffer_records() if r.get("subtype") == AI_MODEL_INVOCATION) + assert call["custom_metadata"]["model"] == "claude-opus-5" + assert call["custom_metadata"]["llm_usage"] == {"input_tokens": 1200, "output_tokens": 340} + + +def test_error_status_marks_the_tool_failed(config, buffer_records): + """An ERROR span status becomes an errored task with its description as stderr.""" + bad = span(**{"gen_ai.tool.name": "deploy"}) + bad["status"] = {"status_code": "ERROR", "description": "permission denied"} + otel.ingest_spans([bad], config) + + tool = next(r for r in buffer_records() if r.get("subtype") == AGENT_TOOL) + assert tool["status"] == "ERROR" + assert tool["stderr"] == "permission denied" + + +def test_non_genai_spans_are_ignored(config, buffer_records): + """An HTTP client span is not provenance.""" + assert otel.ingest_spans([span(**{"http.method": "GET", "http.url": "https://x"})], config) == 0 + assert buffer_records() == [] + + +def test_spans_without_a_conversation_id_are_ignored(config): + """A span with no conversation id cannot be grouped and is skipped.""" + orphan = {"name": "tool", "attributes": {"gen_ai.tool.name": "x"}} + assert otel.ingest_spans([orphan], config) == 0 + + +def test_spans_group_into_one_session(config, buffer_records): + """Spans sharing a conversation id land in one workflow.""" + otel.ingest_spans( + [ + span(**{"gen_ai.operation.name": "chat", "gen_ai.request.model": "m"}), + span(**{"gen_ai.tool.name": "read_file", "gen_ai.tool.call.id": "c1"}), + span(**{"gen_ai.tool.name": "write_file", "gen_ai.tool.call.id": "c2"}), + ], + config, + ) + workflows = [r for r in buffer_records() if r.get("type") == "workflow"] + assert len(workflows) == 1 + assert len({r["workflow_id"] for r in buffer_records() if r.get("type") == "task"}) == 1 + + +def test_mixed_gen_ai_system_does_not_split_the_session(config, buffer_records): + """Spans of one conversation land in ONE workflow even when only some set `gen_ai.system`.""" + otel.ingest_spans( + [ + span(**{CONVERSATION: "conv-mixed", "gen_ai.operation.name": "chat", "gen_ai.request.model": "m"}), + span( + **{ + CONVERSATION: "conv-mixed", + "gen_ai.system": "openai", + "gen_ai.tool.name": "grep", + "gen_ai.tool.call.id": "c1", + } + ), + ], + config, + ) + + records = buffer_records() + assert len([r for r in records if r.get("type") == "workflow"]) == 1 + assert len({r["workflow_id"] for r in records if r.get("type") == "task"}) == 1 + # The provider name is still recorded, as a one-time lifecycle event. + notice = next(r for r in records if r.get("subtype") == HARNESS_EVENT) + assert notice["used"] == {"trigger": "gen_ai.system", "message": "openai"} + + +def test_different_conversations_stay_separate_sessions(config, buffer_records): + """Two conversation ids still yield two workflows.""" + otel.ingest_spans( + [ + span(**{CONVERSATION: "conv-a", "gen_ai.tool.name": "alpha"}), + span(**{CONVERSATION: "conv-b", "gen_ai.tool.name": "beta"}), + ], + config, + ) + workflows = {r["workflow_id"] for r in buffer_records() if r.get("type") == "workflow"} + assert len(workflows) == 2 + + +def test_conflicting_gen_ai_system_first_value_wins(config, buffer_records): + """Document the chosen policy: first-wins. + + The first non-empty `gen_ai.system` a conversation shows is the one + recorded; later, different values neither re-record nor split the session. + """ + otel.ingest_spans( + [ + span( + **{ + CONVERSATION: "conv-conflict", + "gen_ai.system": "openai", + "gen_ai.tool.name": "t1", + "gen_ai.tool.call.id": "c1", + } + ), + span( + **{ + CONVERSATION: "conv-conflict", + "gen_ai.system": "anthropic", + "gen_ai.tool.name": "t2", + "gen_ai.tool.call.id": "c2", + } + ), + ], + config, + ) + + records = buffer_records() + assert len([r for r in records if r.get("type") == "workflow"]) == 1 + notices = [r for r in records if r.get("subtype") == HARNESS_EVENT] + assert len(notices) == 1 + assert notices[0]["used"]["message"] == "openai" + + +def test_ingest_jsonl_file(config, tmp_path, buffer_records): + """Spans written as JSON lines are ingested from disk.""" + path = tmp_path / "spans.jsonl" + path.write_text( + "\n".join( + json.dumps(span(**{"gen_ai.tool.name": name, "gen_ai.tool.call.id": name})) for name in ("alpha", "beta") + ), + encoding="utf-8", + ) + assert otel.ingest_file(path, config) == 2 + assert {r["activity_id"] for r in buffer_records() if r.get("subtype") == AGENT_TOOL} == {"alpha", "beta"} + + +def test_ingest_json_array_file(config, tmp_path): + """Spans written as one JSON array are ingested from disk.""" + path = tmp_path / "spans.json" + path.write_text(json.dumps([span(**{"gen_ai.tool.name": "alpha"})]), encoding="utf-8") + assert otel.ingest_file(path, config) == 1 + + +def test_ingest_otlp_envelope(config, buffer_records): + """Collectors emit OTLP, whose attributes are a list of typed key-values.""" + envelope = { + "resourceSpans": [ + { + "scopeSpans": [ + { + "spans": [ + { + "name": "tool", + "attributes": [ + {"key": CONVERSATION, "value": {"stringValue": "conv-9"}}, + {"key": "gen_ai.tool.name", "value": {"stringValue": "grep"}}, + {"key": "gen_ai.usage.input_tokens", "value": {"intValue": 42}}, + ], + "startTimeUnixNano": 1_700_000_000_000_000_000, + "endTimeUnixNano": 1_700_000_000_500_000_000, + } + ] + } + ] + } + ] + } + assert otel.ingest_spans(otel._flatten([envelope]), config) == 1 + assert next(r for r in buffer_records() if r.get("subtype") == AGENT_TOOL)["activity_id"] == "grep" + + +def test_exporter_records_live_spans(config, buffer_records): + """Drive the exporter through a real tracer provider, not a stub.""" + pytest.importorskip("opentelemetry.sdk") + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(otel.FlowceptSpanExporter(config))) + tracer = provider.get_tracer("test") + + with tracer.start_as_current_span("chat") as s: + s.set_attribute(CONVERSATION, "live-1") + s.set_attribute("gen_ai.operation.name", "chat") + s.set_attribute("gen_ai.request.model", "claude-opus-5") + with tracer.start_as_current_span("tool") as s: + s.set_attribute(CONVERSATION, "live-1") + s.set_attribute("gen_ai.tool.name", "run_tests") + provider.shutdown() + + records = buffer_records() + assert any(r.get("subtype") == AI_MODEL_INVOCATION for r in records) + assert any(r.get("activity_id") == "run_tests" for r in records) + + +def test_exporter_survives_a_malformed_span(config): + """A bad span must not take down the exporting process.""" + exporter = otel.FlowceptSpanExporter(config) + exporter.export([object()]) # no attributes, no context, no status diff --git a/tests/harness/test_plugin_assets.py b/tests/harness/test_plugin_assets.py new file mode 100644 index 00000000..c4923f2e --- /dev/null +++ b/tests/harness/test_plugin_assets.py @@ -0,0 +1,171 @@ +"""Static checks on the Claude Code plugin's shipped assets. + +Everything the plugin ships is declarative (JSON manifests, SKILL.md files, +shell shims), so these tests validate the files themselves: manifests parse, +referenced scripts exist and are executable, skills carry valid frontmatter +and cite only paths that exist in this repository, and the optional +auto-report hook stays silent unless explicitly enabled. Stdlib only. +""" + +from __future__ import annotations + +import json +import os +import re +import subprocess +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +PLUGIN_ROOT = REPO_ROOT / "plugins" / "flowcept" +SKILLS = PLUGIN_ROOT / "skills" + +_PATH_TOKEN = re.compile(r"`([^`]+)`") +_REPO_PATH = re.compile(r"^(?:src|tests|examples|docs|plugins)/[\w./-]+$|^pyproject\.toml$") + + +def _load_json(path: Path) -> dict: + """Parse a JSON file, failing the test with a readable message on error.""" + assert path.is_file(), f"missing {path}" + return json.loads(path.read_text(encoding="utf-8")) + + +def _frontmatter(skill_md: Path) -> dict[str, str]: + """Parse the ``--- ... ---`` YAML-ish frontmatter of a SKILL.md into a dict.""" + lines = skill_md.read_text(encoding="utf-8").splitlines() + assert lines and lines[0].strip() == "---", f"{skill_md} does not start with frontmatter" + fields: dict[str, str] = {} + for line in lines[1:]: + if line.strip() == "---": + return fields + key, sep, value = line.partition(":") + if sep: + fields[key.strip()] = value.strip() + raise AssertionError(f"{skill_md} frontmatter never closes") + + +def _cited_repo_paths(skill_md: Path) -> list[str]: + """Return every repo path cited in a skill.""" + out: list[str] = [] + for line in skill_md.read_text(encoding="utf-8").splitlines(): + for token in _PATH_TOKEN.findall(line): + if _REPO_PATH.match(token): + out.append(token) + return out + + +# -- hooks --------------------------------------------------------------------- + + +def test_hooks_json_parses_and_scripts_are_executable(): + """Every script referenced by hooks.json exists and is executable.""" + hooks = _load_json(PLUGIN_ROOT / "hooks" / "hooks.json")["hooks"] + commands = [h["command"] for entries in hooks.values() for entry in entries for h in entry["hooks"]] + assert commands, "hooks.json declares no commands" + for command in commands: + match = re.search(r"\$\{CLAUDE_PLUGIN_ROOT\}(/[^\"]+)", command) + assert match, f"unrecognized hook command: {command}" + script = PLUGIN_ROOT / match.group(1).lstrip("/") + assert script.is_file(), f"missing script for hook command: {command}" + assert os.access(script, os.X_OK), f"not executable: {script}" + + +def test_hooks_json_capture_session_end_entry_untouched(): + """The original capture entry on SessionEnd is still first and unchanged.""" + hooks = _load_json(PLUGIN_ROOT / "hooks" / "hooks.json")["hooks"] + first = hooks["SessionEnd"][0]["hooks"][0] + assert first["command"] == '"${CLAUDE_PLUGIN_ROOT}/scripts/hook.sh" SessionEnd' + assert first["timeout"] == 10 + + +def test_hooks_json_autoreport_entry_added(): + """A second SessionEnd entry invokes autoreport.sh with a generous timeout.""" + hooks = _load_json(PLUGIN_ROOT / "hooks" / "hooks.json")["hooks"] + assert len(hooks["SessionEnd"]) == 2 + auto = hooks["SessionEnd"][1]["hooks"][0] + assert "autoreport.sh" in auto["command"] + assert auto["timeout"] == 30 + + +# -- manifests ------------------------------------------------------------------- + + +def test_plugin_json_fields(): + """plugin.json parses, is at 0.2.0, and mentions the analysis surface.""" + manifest = _load_json(PLUGIN_ROOT / ".claude-plugin" / "plugin.json") + assert manifest["name"] == "flowcept" + assert manifest["version"] == "0.2.0" + assert "MCP" in manifest["description"] + assert "analy" in manifest["description"].lower() + + +def test_mcp_json_declares_provenance_server(): + """.mcp.json declares the flowcept-provenance stdio server via a real script.""" + manifest = _load_json(PLUGIN_ROOT / ".mcp.json") + server = manifest["mcpServers"]["flowcept-provenance"] + script = PLUGIN_ROOT / server["command"].replace("${CLAUDE_PLUGIN_ROOT}/", "") + assert script.is_file() and os.access(script, os.X_OK) + text = script.read_text(encoding="utf-8") + assert "flowcept.agents.harness.mcp_server" in text + assert "--transport stdio" in text + + +def test_mcp_server_module_runs_as_main(): + """The module the launcher invokes with -m has a __main__ guard.""" + module = REPO_ROOT / "src" / "flowcept" / "agents" / "harness" / "mcp_server.py" + assert 'if __name__ == "__main__"' in module.read_text(encoding="utf-8") + + +def test_marketplace_json_fields(): + """marketplace.json parses, is at 0.2.0, and points at an existing plugin dir.""" + manifest = _load_json(REPO_ROOT / ".claude-plugin" / "marketplace.json") + assert manifest["metadata"]["version"] == "0.2.0" + (entry,) = [p for p in manifest["plugins"] if p["name"] == "flowcept"] + assert (REPO_ROOT / entry["source"]).is_dir() + assert "analy" in entry["description"].lower() + + +# -- skills ---------------------------------------------------------------------- + + +@pytest.mark.parametrize("skill", ["session-provenance", "prov-analysis", "write-flowcept-plugin"]) +def test_skill_frontmatter(skill): + """Each skill has frontmatter whose name matches its directory.""" + fields = _frontmatter(SKILLS / skill / "SKILL.md") + assert fields.get("name") == skill + assert len(fields.get("description", "")) > 40 + + +@pytest.mark.parametrize("skill", ["prov-analysis", "write-flowcept-plugin"]) +def test_skill_cited_paths_exist(skill): + """Every repo path a new skill cites exists.""" + cited = _cited_repo_paths(SKILLS / skill / "SKILL.md") + assert cited, f"{skill} cites no repo paths" + for path in cited: + assert (REPO_ROOT / path).exists(), f"{skill} cites missing path: {path}" + + +def test_session_provenance_skill_intact(): + """The capture-side skill still documents the buffer location and record shapes.""" + text = (SKILLS / "session-provenance" / "SKILL.md").read_text(encoding="utf-8") + for needle in ("buffers/.jsonl", "agent_tool", "ai_model_invocation", "flowcept-harness flush"): + assert needle in text + + +# -- autoreport behavior ----------------------------------------------------------- + + +def test_autoreport_silent_noop_when_disabled(): + """With FLOWCEPT_HARNESS_AUTOREPORT unset, autoreport.sh exits 0 with no stdout.""" + env = {k: v for k, v in os.environ.items() if k != "FLOWCEPT_HARNESS_AUTOREPORT"} + result = subprocess.run( + [str(PLUGIN_ROOT / "scripts" / "autoreport.sh")], + input=json.dumps({"hook_event_name": "SessionEnd", "session_id": "fake"}).encode(), + capture_output=True, + env=env, + timeout=30, + check=False, + ) + assert result.returncode == 0 + assert result.stdout == b"" diff --git a/tests/harness/test_sdk.py b/tests/harness/test_sdk.py new file mode 100644 index 00000000..d2e5778e --- /dev/null +++ b/tests/harness/test_sdk.py @@ -0,0 +1,471 @@ +"""Tests for the in-process SDK wrappers. + +The three SDKs are duck-typed, so these drive each wrapper with objects shaped +like the real payloads rather than requiring the SDKs to be installed. Where an +SDK *is* installed, a further test drives the real thing. +""" + +from __future__ import annotations + +import pytest + +from flowcept.agents.harness import SessionTracer +from flowcept.agents.claude_agent_sdk.claude_agent_sdk_plugin import ClaudeAgentTracer +from flowcept.agents.langchain.langchain_plugin import FlowceptCallbackHandler +from flowcept.agents.openai_agents.openai_agents_plugin import FlowceptTraceProcessor +from flowcept.agents.harness.vocab import AGENT_TOOL, AI_MODEL_INVOCATION, SUBAGENT_SESSION + + +def tasks(records, subtype=None): + return [r for r in records if r.get("type") == "task" and (subtype is None or r.get("subtype") == subtype)] + + +def workflows(records, subtype=None): + return [r for r in records if r.get("type") == "workflow" and (subtype is None or r.get("subtype") == subtype)] + + +# -- SessionTracer ------------------------------------------------------------ + + +def test_tracer_records_a_whole_run(config, buffer_records): + with SessionTracer("my_agent", "run-1", config=config, model="claude-opus-5") as tracer: + tracer.prompt("summarize the repo") + call = tracer.tool_start("read_file", {"path": "README.md"}) + tracer.tool_end(call, tool_response={"bytes": 4096}) + tracer.turn_end("done", usage={"input_tokens": 900, "output_tokens": 30}) + + records = buffer_records() + turn = tasks(records, AI_MODEL_INVOCATION)[0] + tool = tasks(records, AGENT_TOOL)[0] + assert turn["used"]["prompt"] == "summarize the repo" + assert turn["custom_metadata"]["llm_usage"]["input_tokens"] == 900 + assert tool["activity_id"] == "read_file" + assert tool["generated"]["bytes"] == 4096 + # The tool hangs off the turn, which is the edge the whole thing is for. + assert tool["parent_task_id"] == turn["task_id"] + + session = workflows(records)[-1] + assert session["status"] == "FINISHED" + assert session["generated"] == {"turns": 1, "tool_calls": 1} + + +def test_tool_context_manager_records_a_raised_exception(config, buffer_records): + tracer = SessionTracer("my_agent", "run-2", config=config) + with pytest.raises(ValueError): + with tracer.tool("run_tests", {"suite": "unit"}) as call: + call.result({"ignored": True}) + raise ValueError("boom") + + tool = tasks(buffer_records(), AGENT_TOOL)[0] + assert tool["status"] == "ERROR" + assert tool["stderr"] == "ValueError: boom" + + +def test_tracer_end_is_idempotent(config, buffer_records): + tracer = SessionTracer("my_agent", "run-3", config=config) + tracer.prompt("hi") + tracer.end() + tracer.end() + assert len([w for w in workflows(buffer_records()) if w.get("status") == "FINISHED"]) == 1 + + +def test_an_interrupted_run_still_closes_its_turn(config, buffer_records): + """A crash mid-turn must leave attributable provenance, not nothing.""" + with pytest.raises(RuntimeError): + with SessionTracer("my_agent", "run-4", config=config) as tracer: + tracer.prompt("long job") + raise RuntimeError("killed") + + turn = tasks(buffer_records(), AI_MODEL_INVOCATION)[0] + assert turn["custom_metadata"]["close_reason"] == "session_ended" + assert workflows(buffer_records())[-1]["custom_metadata"]["end_reason"] == "error" + + +# -- Claude Agent SDK --------------------------------------------------------- + + +class Block: + def __init__(self, **fields): + self.__dict__.update(fields) + + +class AssistantMessage: + def __init__(self, content, model="claude-opus-5"): + self.content = content + self.model = model + + +class UserMessage: + def __init__(self, content): + self.content = content + + +class ResultMessage: + def __init__(self, result=None, usage=None, is_error=False, session_id=None): + self.result = result + self.usage = usage + self.is_error = is_error + self.session_id = session_id + self.num_turns = 1 + self.total_cost_usd = 0.01 + + +class SystemMessage: + def __init__(self, subtype, data=None): + self.subtype = subtype + self.data = data or {} + + +def test_claude_agent_stream_becomes_provenance(config, buffer_records): + with ClaudeAgentTracer(config=config, prompt="fix the failing test") as tracer: + tracer.handle( + AssistantMessage( + [ + Block(text="Let me look."), + Block(id="tu_1", name="Read", input={"file_path": "test_x.py"}), + ] + ) + ) + tracer.handle(UserMessage([Block(tool_use_id="tu_1", content="def test_x(): ...", is_error=False)])) + tracer.handle(ResultMessage(result="fixed", usage={"input_tokens": 1200, "output_tokens": 80})) + + records = buffer_records() + turn = tasks(records, AI_MODEL_INVOCATION)[0] + assert turn["used"]["prompt"] == "fix the failing test" + assert turn["generated"]["response"] == "fixed" + # Cost travels with usage; both are on the turn, which is where a query for + # "what did this session cost" will look. + assert turn["custom_metadata"]["llm_usage"]["total_cost_usd"] == 0.01 + assert turn["custom_metadata"]["llm_usage"]["input_tokens"] == 1200 + + tool = tasks(records, AGENT_TOOL)[0] + assert tool["activity_id"] == "Read" + assert tool["used"]["file_path"] == "test_x.py" + assert tool["status"] == "FINISHED" + + +def test_claude_agent_adopts_the_sdk_session_id(config, buffer_records): + """The init message's id is taken, so a resumed run reuses the workflow.""" + tracer = ClaudeAgentTracer(config=config) + tracer.handle(SystemMessage("init", {"session_id": "sdk-abc", "model": "claude-opus-5"})) + tracer.handle(AssistantMessage([Block(text="hi")])) + tracer.close() + + assert tracer.tracer.session_id == "sdk-abc" + + from flowcept.agents.harness import ids + + expected = ids.workflow_id_for("claude_agent_sdk", "sdk-abc") + assert {r["workflow_id"] for r in buffer_records()} == {expected} + + +def test_claude_agent_keeps_an_explicit_session_id(config): + tracer = ClaudeAgentTracer("mine", config=config) + tracer.handle(SystemMessage("init", {"session_id": "sdk-abc"})) + assert tracer.tracer.session_id == "mine" + + +def test_claude_agent_task_tool_opens_a_subagent_workflow(config, buffer_records): + with ClaudeAgentTracer(config=config, prompt="explore") as tracer: + tracer.handle( + AssistantMessage( + [Block(id="tu_task", name="Task", input={"subagent_type": "Explore", "prompt": "find the tests"})] + ) + ) + tracer.handle(UserMessage([Block(tool_use_id="tu_task", content="found 3", is_error=False)])) + tracer.handle(ResultMessage(result="ok")) + + subagent = workflows(buffer_records(), SUBAGENT_SESSION) + assert len(subagent) == 1 # opened and closed, the open record superseded + assert subagent[0]["used"] == {"agent_type": "Explore", "prompt": "find the tests"} + assert subagent[0]["generated"]["response"] == "found 3" + assert subagent[0]["status"] == "FINISHED" + + +def test_claude_agent_failed_tool_is_an_error(config, buffer_records): + with ClaudeAgentTracer(config=config, prompt="p") as tracer: + tracer.handle(AssistantMessage([Block(id="tu_1", name="Bash", input={"command": "false"})])) + tracer.handle(UserMessage([Block(tool_use_id="tu_1", content="exit status 1", is_error=True)])) + + tool = tasks(buffer_records(), AGENT_TOOL)[0] + assert tool["status"] == "ERROR" + assert tool["stderr"] == "exit status 1" + + +def test_claude_agent_closes_a_tool_that_never_returned(config, buffer_records): + tracer = ClaudeAgentTracer(config=config, prompt="p") + tracer.handle(AssistantMessage([Block(id="tu_1", name="Bash", input={"command": "sleep 999"})])) + tracer.close(error="interrupted") + + tool = tasks(buffer_records(), AGENT_TOOL)[0] + assert tool["status"] == "ERROR" + assert tool["stderr"] == "never returned a result" + + +def test_claude_agent_records_nothing_for_an_empty_run(config, buffer_records): + ClaudeAgentTracer(config=config).close() + assert buffer_records() == [] + + +def test_claude_agent_compaction_is_a_lifecycle_event(config, buffer_records): + tracer = ClaudeAgentTracer(config=config, prompt="p") + tracer.handle(AssistantMessage([Block(text="working")])) + tracer.handle(SystemMessage("compact_boundary")) + tracer.close() + + assert [t["activity_id"] for t in tasks(buffer_records(), "harness_event")] == ["compact"] + + +# -- OpenAI Agents SDK -------------------------------------------------------- + + +class Trace: + def __init__(self, trace_id="tr_1", group_id=None, name="run"): + self.trace_id = trace_id + self.group_id = group_id + self.name = name + + +class Span: + def __init__(self, span_id, data, *, trace_id="tr_1", parent_id=None, error=None): + self.span_id = span_id + self.trace_id = trace_id + self.parent_id = parent_id + self.span_data = data + self.error = error + self.started_at = "2026-08-19T10:00:00+00:00" + self.ended_at = "2026-08-19T10:00:02+00:00" + + +class SpanData: + def __init__(self, type, **fields): + self.type = type + self.__dict__.update(fields) + + +def test_openai_agents_trace_becomes_a_session(config, buffer_records): + processor = FlowceptTraceProcessor(config) + trace = Trace(group_id="thread-9") + processor.on_trace_start(trace) + + generation = Span("sp_1", SpanData("generation", model="gpt-5", input="hi", output="hello", + usage={"input_tokens": 10, "output_tokens": 3})) + processor.on_span_start(generation) + processor.on_span_end(generation) + + function = Span("sp_2", SpanData("function", name="get_weather", input='{"city": "Paris"}', output="18C")) + processor.on_span_start(function) + processor.on_span_end(function) + + processor.on_trace_end(trace) + + records = buffer_records() + call = next(t for t in tasks(records, AI_MODEL_INVOCATION) if t["activity_id"] == "llm_interaction") + assert call["custom_metadata"]["model"] == "gpt-5" + assert call["custom_metadata"]["llm_usage"]["output_tokens"] == 3 + # The span reported its own duration; it must survive. + assert call["ended_at"] - call["started_at"] == pytest.approx(2.0, abs=0.01) + + tool = tasks(records, AGENT_TOOL)[0] + assert tool["activity_id"] == "get_weather" + assert tool["generated"]["value"] == "18C" + + session = workflows(records)[-1] + assert session["status"] == "FINISHED" + + +def test_openai_agents_nested_agent_becomes_a_subagent(config, buffer_records): + processor = FlowceptTraceProcessor(config) + trace = Trace() + processor.on_trace_start(trace) + + root = Span("sp_root", SpanData("agent", name="Triage", model="gpt-5")) + processor.on_span_start(root) + nested = Span("sp_sub", SpanData("agent", name="Researcher", output="found it"), parent_id="sp_root") + processor.on_span_start(nested) + tool = Span("sp_tool", SpanData("function", name="search", output="hits"), parent_id="sp_sub") + processor.on_span_start(tool) + processor.on_span_end(tool) + processor.on_span_end(nested) + processor.on_span_end(root) + processor.on_trace_end(trace) + + records = buffer_records() + subagents = workflows(records, SUBAGENT_SESSION) + assert [w["name"] for w in subagents] == ["subagent:Researcher"] + # A root agent is the session, so it must not also become a subagent. + assert not any("Triage" in w["name"] for w in subagents) + # The nested agent's tool belongs to the nested agent's workflow. + assert tasks(records, AGENT_TOOL)[0]["workflow_id"] == subagents[0]["workflow_id"] + + +def test_openai_agents_span_error_marks_the_tool_failed(config, buffer_records): + processor = FlowceptTraceProcessor(config) + trace = Trace() + processor.on_trace_start(trace) + span = Span("sp_1", SpanData("function", name="deploy"), error={"message": "denied", "data": "no creds"}) + processor.on_span_start(span) + processor.on_span_end(span) + processor.on_trace_end(trace) + + tool = tasks(buffer_records(), AGENT_TOOL)[0] + assert tool["status"] == "ERROR" + assert tool["stderr"] == "denied: no creds" + + +def test_openai_agents_tripped_guardrail_is_an_error(config, buffer_records): + processor = FlowceptTraceProcessor(config) + trace = Trace() + processor.on_trace_start(trace) + span = Span("sp_1", SpanData("guardrail", name="no_pii", triggered=True)) + processor.on_span_start(span) + processor.on_span_end(span) + processor.on_trace_end(trace) + + tool = tasks(buffer_records(), AGENT_TOOL)[0] + assert tool["activity_id"] == "no_pii" + assert tool["stderr"] == "guardrail triggered" + + +def test_openai_agents_ignores_spans_from_an_unknown_trace(config, buffer_records): + """Spans can outlive their trace; they must not create a phantom session.""" + processor = FlowceptTraceProcessor(config) + processor.on_span_end(Span("sp_1", SpanData("function", name="x"))) + assert buffer_records() == [] + + +def test_openai_agents_shutdown_closes_open_traces(config, buffer_records): + processor = FlowceptTraceProcessor(config) + processor.on_trace_start(Trace()) + processor.shutdown() + assert workflows(buffer_records())[-1]["custom_metadata"]["end_reason"] == "shutdown" + + +# -- LangChain / LangGraph ---------------------------------------------------- + + +class Generation: + def __init__(self, text): + self.text = text + self.message = None + + +class LLMResult: + def __init__(self, text, llm_output=None): + self.generations = [[Generation(text)]] + self.llm_output = llm_output + + +def test_langchain_chain_run_becomes_a_turn(config, buffer_records): + handler = FlowceptCallbackHandler("thread-1", config=config) + handler.on_chain_start({"name": "AgentExecutor"}, {"input": "what is 2+2"}, run_id="r0") + handler.on_llm_start( + {"name": "ChatOpenAI"}, ["what is 2+2"], run_id="r1", parent_run_id="r0", + invocation_params={"model": "gpt-5"}, + ) + handler.on_llm_end(LLMResult("4", {"token_usage": {"prompt_tokens": 12, "completion_tokens": 1}}), run_id="r1") + handler.on_chain_end({"output": "4"}, run_id="r0") + handler.close() + + records = buffer_records() + turn = next(t for t in tasks(records, AI_MODEL_INVOCATION) if t["activity_id"] == "agent_turn") + assert turn["used"]["prompt"] == "what is 2+2" + assert turn["generated"]["response"] == "4" + + call = next(t for t in tasks(records, AI_MODEL_INVOCATION) if t["activity_id"] == "llm_interaction") + assert call["custom_metadata"]["model"] == "gpt-5" + assert call["custom_metadata"]["llm_usage"]["prompt_tokens"] == 12 + # The call nests under the turn, not beside it. + assert call["parent_task_id"] == turn["task_id"] + + +def test_langchain_tool_run_becomes_a_tool_task(config, buffer_records): + handler = FlowceptCallbackHandler("thread-2", config=config) + handler.on_chain_start(None, {"input": "weather?"}, run_id="r0") + handler.on_tool_start({"name": "get_weather"}, '{"city": "Paris"}', run_id="r1", + parent_run_id="r0", inputs={"city": "Paris"}) + handler.on_tool_end("18C", run_id="r1") + handler.on_chain_end({"output": "18C"}, run_id="r0") + handler.close() + + tool = tasks(buffer_records(), AGENT_TOOL)[0] + assert tool["activity_id"] == "get_weather" + assert tool["used"] == {"city": "Paris"} + assert tool["generated"]["value"] == "18C" + + +def test_langchain_nested_chains_are_not_recorded(config, buffer_records): + """A LangGraph run emits a chain per node; only the outer one is a turn.""" + handler = FlowceptCallbackHandler("thread-3", config=config) + handler.on_chain_start(None, {"messages": ["go"]}, run_id="r0") + for node in ("agent", "tools", "agent"): + handler.on_chain_start({"name": node}, {}, run_id=f"n_{node}", parent_run_id="r0") + handler.on_chain_end({}, run_id=f"n_{node}") + handler.on_chain_end({"messages": ["done"]}, run_id="r0") + handler.close() + + assert len(tasks(buffer_records(), AI_MODEL_INVOCATION)) == 1 + + +def test_langchain_tool_error_is_recorded(config, buffer_records): + handler = FlowceptCallbackHandler("thread-4", config=config) + handler.on_tool_start({"name": "deploy"}, "", run_id="r1") + handler.on_tool_error(RuntimeError("denied"), run_id="r1") + handler.close() + + tool = tasks(buffer_records(), AGENT_TOOL)[0] + assert tool["status"] == "ERROR" + assert tool["stderr"] == "RuntimeError: denied" + + +def test_langchain_bare_model_call_is_its_own_turn(config, buffer_records): + """No chain around it: the model call is the whole turn.""" + handler = FlowceptCallbackHandler("thread-5", config=config) + handler.on_llm_start({"id": ["langchain", "ChatAnthropic"]}, ["hello"], run_id="r1") + handler.on_llm_end(LLMResult("hi there"), run_id="r1") + handler.close() + + turn = next(t for t in tasks(buffer_records(), AI_MODEL_INVOCATION) if t["activity_id"] == "agent_turn") + assert turn["used"]["prompt"] == "hello" + assert turn["generated"]["response"] == "hi there" + assert turn["status"] == "FINISHED" + + +def test_langchain_chain_error_marks_the_turn_failed(config, buffer_records): + handler = FlowceptCallbackHandler("thread-6", config=config) + handler.on_chain_start(None, {"input": "x"}, run_id="r0") + handler.on_chain_error(ValueError("bad graph"), run_id="r0") + handler.close() + + turn = tasks(buffer_records(), AI_MODEL_INVOCATION)[0] + assert turn["status"] == "ERROR" + assert turn["stderr"] == "ValueError: bad graph" + + +def test_langchain_chat_model_start_flattens_messages(config, buffer_records): + class Message: + def __init__(self, type, content): + self.type = type + self.content = content + + handler = FlowceptCallbackHandler("thread-7", config=config) + handler.on_chat_model_start( + {"name": "ChatOpenAI"}, + [[Message("system", "be brief"), Message("human", "hi")]], + run_id="r1", + invocation_params={"model": "gpt-5"}, + ) + handler.on_llm_end(LLMResult("hello"), run_id="r1") + handler.close() + + call = next(t for t in tasks(buffer_records(), AI_MODEL_INVOCATION) if t["activity_id"] == "llm_interaction") + assert call["used"]["prompt"] == "system: be brief\nhuman: hi" + + +def test_langchain_handler_exposes_the_manager_contract(config): + """langchain's callback manager reads these off the handler by name.""" + handler = FlowceptCallbackHandler(config=config) + for attribute in ("ignore_llm", "ignore_chain", "ignore_agent", "ignore_retriever", + "ignore_chat_model", "ignore_retry", "ignore_custom_event", + "raise_error", "run_inline"): + assert isinstance(getattr(handler, attribute), bool)