Orchestrata Ultimate › Plans › context-injection-evolution-plan

Context Injection Evolution Plan

Source: plans/context-injection-evolution-plan.md at 8a2cb42ad · last changed 2026-07-27 · first committed 2026-07-13 · 367 lines.

Prose pending.

Plan text

Context Injection Evolution Plan

Analysis of how the project currently assembles an agent's prompt context, plus a concrete plan to make far better use of the AST, Neo4j knowledge graph, and session/plan/memory context — while keeping prompts inside a bounded, sliding context window.

1. Current state and gaps

1.1 What a volk receives today

The pipeline uses a layered prepend model: each context source is rendered as markdown and glued in front of the user message (src/orchestrata_ultimate/llm/client.py).

The live injector already does several things well:

• It is subject-gated (implement, fix_issues, review_and_test, design_review, scope_review, ux_design).
• It returns a token-capped prompt (GraphContext.to_prompt(max_context_tokens=4000)).
• It includes semantic layer facts: requirements via TRACES_TO and decisions via APPLIES_TO.
• The newer ContextPack path is deterministic, editable, and budgeted (packer.py).

1.2 Key gaps

1. No cross-layer budget. GraphContext, ContextPack, memory recall, and the user message each have their own limits, but no component tracks the total prompt size.
2. Session/plan state is missing. Volks have no compact view of the active goal, pending work items, handoff notes, or the current plan step.
3. Artifact feedback is manual. Prior-stage outputs rely on each volk copying the right payload keys.
4. Memory is inconsistently wired. Recall only works when a MemoryEngine is registered for the thread.
5. AST is underused. The AST engine (scripts/serialize_codebase_ast.py) produces branches, complexity, call graphs, untyped/dead-param reports, and skeletons, but none of it reaches volk prompts.
6. Graph traversal is shallow. The live injector uses basename matching and a small fixed Cypher expansion; it does not exploit call-graph directionality, centrality, or semantic search.
7. No feedback loop. GraphContextValidator checks whether the response references injected entities, but a failed validation does not reformulate the injected context.
8. No on-demand retrieval. Everything is prepended. Volks cannot ask for more context when they need it.

2. Leveraging the AST

2.1 What the AST engine already provides

scripts/serialize_codebase_ast.py (and its FastAPI wrapper scripts/ast_api.py) reads live Python source and produces structured, prompt-usable output:

The AST is also synced into Neo4j by src/orchestrata_ultimate/knowledge_graph/sync.py as 12 node labels (Project, Package, Module, Class, Function, Method, Parameter, Requirement, ImportNode, DecoratorNode, Spec, RequirementItem) and 18 relationship types (DEFINES, HAS_METHOD, IMPORTS, CALLS, INHERITS_FROM, TRACES_TO, TESTED_BY, etc.).

2.2 Proposed AST-driven context features

1. Branch-aware prompts. Before modifying a function, run serialize codebase_ast branches --name <symbol> and inject the decision points. The agent must explicitly address each if/except path.
2. Complexity-first dispatch. For fix_issues or review_and_test, inject the top-N complexity hotspots so the agent prioritizes high-risk functions.
3. Call-graph neighborhood. For an implement or fix_issues packet, fetch call-graph --name <symbol> callers and callees and add them to the KG context. This reduces duplicate-code risk better than file-level imports.
4. Skeleton contract workflow. Use gen-skeleton <contract.mmd> to emit the exact target shape, hand it to the agent as the implementation contract, then run skeleton-diff to detect structural escapes.
5. Trace-lint preflight. Before a task claiming to satisfy requirements, run trace-lint and inject any dangling/placeholder findings. The agent must resolve them before marking work done.
6. Untested-branch scaffolding. Combine branches + gen-tests to generate branch-level tests and require the agent to make them pass.
7. API-on-demand. Keep scripts/ast_api.py running so volks can call /symbols, /branches, /refs, /complexity dynamically instead of parsing source themselves.

3. Leveraging the Knowledge Graph

3.1 What the KG already provides

• canonical_schema.py: 51 node labels and canonical identity ({project}:{natural_id}).
• query_engine.py: generic Cypher, multi-hop traversal, provenance lookup, lexical search.
• context_injector.py: file/relationship/test/requirement/decision context.
• context_pack/: deterministic, task-type-templated, ranked, budgeted context packs.
• graph_context_validator.py: checks that responses reference injected public entities.
• relevance.py + prolog/kg_relevance.pl: Prolog-deduced core/related/peripheral tiers.
• semantic_search.py: Memgraph-based vector + keyword search over code chunks.

3.2 Proposed KG-driven improvements

1. Canonical-id lookups. Replace f.path CONTAINS $basename matching in ContextInjector._query_context with exact File.id lookups ({project}:{path}), then expand via DEFINES.
2. Call-graph-biased expansion. Weight CALLS edges higher than IMPORTS; use in-degree centrality to surface hub functions.
3. Requirement closure. Expand TRACES_TO to RequirementItem, then to parent Requirement/Spec, so the full "why" hierarchy is visible.
4. Populate missing sections. GraphContext.recent_changes and similar_patterns are currently empty. Query updated_at/synced_at for recency and shared INHERITS_FROM bases for patterns.
5. Semantic search bridge. Map top Memgraph vector-search chunks back to Neo4j File/Class/Function nodes by (file_path, start_line, end_line) and feed them as extra root entities into analyze_request.
6. Validation feedback loop. Make GraphContextValidator node-aware (parse pack provenance or pass node ids). On failure, expand k_hop or include more TRACES_TO/TESTED_BY edges and re-pack.
7. Runtime signal ingestion. Write Message, GatewayEvent, Decision, ToolCall nodes (Tier H labels) into Neo4j so Prolog relevance rules can prioritize recently changed or failing nodes for DEBUG tasks.
8. Unify requirement identity. Merge or relate Requirement (code traces) and RequirementItem (spec traces) so requirement-sharing logic works across both.
9. Cache invalidation. Tie query_engine.py LRU cache to graph synced_at or transaction ID so syncs clear stale entries.

4. Leveraging Session, Plan, and Memory Context

4.1 What session context already provides

src/orchestrata_ultimate/session/models.py defines:

• ProjectState: active session, version, branch, test state, status, current goal, portfolio goals.
• SessionRecord: purpose, objectives, work log, artifacts, choices, achievements, upcoming, handoff.
• HandoffData: decisions, blockers, active work, key files, commands, alerts.
• WorkItem: deferred tasks with priority, state, notes.
• Goal (SPEC-023): statement, criteria, tasks, progress, blockers.

ContextFileManager (session/context.py) already renders context.md via Jinja2 templates in session/templates/. The memory engine (memory/) provides hot (Redis) and cold (Neo4j) stores with a deterministic hash embedder and ContextCompressor. ArtifactStore captures per-stage outputs.

4.2 Proposed session-driven context features

1. Mission block injection. Load ContextData once per pipeline run and render a compact block containing:
   • ProjectState status, branch, test state.
   • Active Goal statement and unmet criteria.
   • Top pending WorkItems.
   • Latest HandoffData (decisions, blockers, key files, alerts).
2. Memory summary + recall. At pipeline start, call MemoryEngine.summarize(context_id=thread_id, max_tokens=512) for a run-level summary; keep per-task recall() for specifics. Prepend both.
3. Artifact feedback loop. After each stage, ArtifactStore.add() already captures outputs. Query artifacts.by_thread(thread_id) and inject a chronological digest (stage, type, file path, 200-char preview).
4. Plan window injection. Parse plans/execution-plan.md (or the active plan document) into steps. Maintain a current_step pointer in ProjectState or Goal. Inject the current step in full plus the next 2 steps as titles.
5. Canonical ArtifactStore. Consolidate kernel/pipeline_artifacts.py::ArtifactStore and orchestrator/artifact_store.py so both runners feed the same artifact context.

5. Target architecture: Unified Context Assembler

Introduce one service that composes all context sources under a single token budget.

┌─────────────────────────────────────────────────────────────┐
│                   UnifiedContextAssembler                   │
│  inputs: payload, subject, task_type, thread_id, volk_alias  │
│  output: context_parts dict + token accounting               │
└─────────────────────────────────────────────────────────────┘
        │
        ▼
┌──────────────┐  ┌──────────────┐  ┌──────────────┐  ┌──────────────┐
│ Session/Goal │  │ Memory       │  │ Artifact     │  │ KG / AST     │
│ Context      │  │ Summary      │  │ Digest       │ │ Context      │
└──────────────┘  └──────────────┘  └──────────────┘  └──────────────┘
        │                │                │                │
        └────────────────┴────────────────┴────────────────┘
                                │
                                ▼
                    ┌───────────────────────┐
                    │   ContextBudget       │
                    │   (enforced layering) │
                    └───────────────────────┘
                                │
                                ▼
                    ┌───────────────────────┐
                    │ effective_message     │
                    │ = join(parts) + user  │
                    └───────────────────────┘

Responsibilities:

• Fetch each source independently.
• Allocate a hard per-layer token budget.
• Evict lowest-relevance items first when a layer overflows.
• Cache rendered packs per (snapshot_id, volk_alias, subject).
• Expose a get_context(entity, depth=1) tool so volks can pull more on demand.
• Emit structured context-usage telemetry.

Integration points:

• kernel/pipeline_kg.py::_maybe_inject_context becomes the orchestration site.
• orchestrator/pipeline_runner.py::_run_stage calls the same assembler (closes the two-runner divergence).
• llm/client.py accepts context_parts directly instead of a single graph_context string, or the assembler is called upstream.

5.1 The voluntary-use problem

Historical experiments show a consistent failure mode: even when instructions tell the model to use provided context, the model does not voluntarily pull or reference the data. It falls back on training-data shortcuts, hallucinates, or ignores the injected material.

This means context injection cannot rely on politeness. The system must force usage:

1. Prepend, don't side-channel. Put context inside the user message tokens so the model must process it (already done in llm/client.py).
2. Make references mandatory. Require the output to list every entity it used: classes, functions, requirements, decisions, prior artifacts. A response with no citations is invalid.
3. Validate citations. GraphContextValidator must check that cited entities were actually injected and that injected critical entities were cited.
4. Use small, unavoidable context. Inject only the entities the task cannot succeed without. A smaller critical set is harder to ignore than a large optional encyclopedia.
5. Structured output. When possible, make the volk fill a template (e.g., "Entities modified:" + "Requirements satisfied:" + "Decisions respected:") rather than write free prose.
6. Mandatory tool calls for expansion. On-demand retrieval should not be optional. The task prompt can require the model to call get_code/get_tests before producing output, and the pipeline rejects output that skipped the tool call.
7. Retry with escalation. First failure: re-prompt with a warning. Second failure: shrink context to the single most relevant entity and ask again. Third failure: escalate to human.
8. Reward signal via provenance. ContextPack already carries provenance comments. Extend this to require the model to echo the node_ids it acted on; score responses by citation coverage.

6. Context-window management strategy

6.1 Budget model

Reserve slots explicitly before the model sees the prompt. Example budget for a 128k-context model with ~8k output:

The assembler rejects or degrades gracefully when the total exceeds the model window.

6.2 Sliding plan window

For long-running plans:

• Current step: full text (goal, actions, verification, expected output).
• Completed steps: one-sentence summary each; drop summaries older than the last 3 steps unless causally linked to the current step.
• Upcoming steps: title + one-line description only.
• Store current_step in ProjectState and update it as the orchestrator advances.

6.3 Eviction policy

Use the existing ranking signals:

1. Drop TIER_PERIPHERAL nodes.
2. Drop TIER_RELATED nodes.
3. Within a tier, drop highest hop_distance first.
4. Never evict: active goal, current-step plan, blockers, failing tests, prior decisions, current diff.

6.4 On-demand retrieval

Instead of prepending everything, expose tools:

• get_code(entity_id, include_callers=false, depth=1) — fetch a class/function and neighbors.
• get_tests(entity_id) — fetch covering tests.
• get_requirements(entity_id) — fetch traced requirements.
• get_artifact(artifact_id) — fetch a prior-stage artifact.
• search_code(query, limit=5) — semantic + lexical search.

This keeps the default prompt small and lets volks expand context only when needed.

7. Phased implementation plan

Phase 0 — Foundations (suggested first milestone)

Goal: establish budget accounting, a single assembler, and close the most urgent gaps.

Phase 1 — AST-driven context

Goal: feed AST-derived facts into the assembler.

Phase 2 — KG enhancements

Goal: richer, canonical, feedback-aware graph context.

Phase 3 — Session, memory, and plan window

Goal: make long-running sessions coherent and plan-aware.

Phase 4 — Optimization and observability

Goal: make context assembly fast, fresh, and measurable.

Phase 5 — Cleanup (after the initial wave)

Goal: pay down the technical debt left behind to keep the initial wave small.

8. Acceptance criteria for the whole effort

1. Every volk call logs context_usage with tokens per layer.
2. Both PipelineRunner implementations use the same UnifiedContextAssembler.
3. Session state (goal, handoff, pending work) is injected for all INJECT_SUBJECTS.
4. Memory recall works in both choreographed and orchestrated runners.
5. Artifact feedback is automatic, not manual payload copying.
6. AST-derived context (branches, complexity, call graph, trace-lint) is available to at least implement, fix_issues, and review_and_test subjects.
7. Graph context uses canonical File.id lookups and includes recent changes / similar patterns.
8. GraphContextValidator failure triggers a one-step context expansion retry.
9. A plan window (current step + compact done/upcoming) is injected for long-running plan execution.
10. Total prompt size never exceeds the configured model window; peripheral nodes are evicted first.
11. Volk outputs cite injected entities; GraphContextValidator rejects outputs missing mandatory citations.

9. Risks and non-goals

Risks

• Models ignore context even when instructed. This is the historical pattern the plan is designed around; mitigation is mandatory citations + validator + retry, not bigger prompts.
• Increasing context size can raise cost and latency even within budget. Measure per-call spend with telemetry.
• Over-eager eviction can drop critical nodes. Keep a "do not evict" list and test it.
• Bridging Memgraph semantic search into Neo4j adds coupling. Consider converging on one store if Memgraph dependency becomes a burden.
• Runtime Tier H ingestion can write a lot of data. Use TTL or periodic compaction.

Non-goals

• This plan does not replace the existing @traces system; it complements it.
• It does not mandate a single LLM model; budget sizes should be model-configurable.
• It does not require removing ContextInjector or ContextPack; it evolves them behind the assembler.

10. Summary of immediate next actions

If you want to start tomorrow, do these in order:

1. Create ContextBudget and UnifiedContextAssembler skeleton.
2. Pick one ArtifactStore and migrate the other runner to it.
3. Inject a mission block from ContextData for INJECT_SUBJECTS.
4. Fix ContextInjector._query_context to use canonical File.id lookups.
5. Add MemoryEngine.summarize and bind memory universally in both runners.
6. Mandate citations in every volk prompt and make GraphContextValidator enforce them.

Those six items close gaps D, E, and C from docs/context-injection.md and address the observed voluntary-use failure mode.

Related

• Plans
• Orchestrata Ultimate — project index
• Workflow
• Tasks