/opsmill/infrahub-mcp

Wiki directory

Browse AI-generated documentation for any repository.

Session Branch Recovery and Reset

Relevant source files

The following files were used as context for generating this wiki page:

Session Branch Recovery and Reset provides a mechanism to handle stale or unwritable branches automatically during MCP tool execution. This feature ensures that if a branch becomes read-only (e.g., after being merged) or is deleted, the system provisions a fresh branch to continue write operations without requiring a server restart. It also introduces a manual tool to explicitly reset or switch the current session's active branch.

Sources: spec.md:12-25, plan.md:10-15

Architecture and State Management

The system manages branch state per MCP session using a WeakKeyDictionary keyed by the session object. This design replaces process-wide caching to provide isolation between concurrent connections and ensures that state is automatically evicted when a session ends.

State Storage Structure

The AppContext holds the in-memory state for session branches and associated locks.

Entity	Type	Description
`_session_branches`	`WeakKeyDictionary[Any, str]`	Maps the session object to the current branch name.
`_session_locks`	`WeakKeyDictionary[Any, asyncio.Lock]`	Provides per-session locking for branch operations.
`_session_locks_guard`	`asyncio.Lock`	Protects the creation of new entries in the lock dictionary.

Sources: data-model.md:12-28, tasks.md:40-45

The following diagram illustrates the flow for resolving or recovering a session branch during a write operation:

Generating architecture diagram...

Sources: data-model.md:58-70, plan.md:10-15

Branch Recovery Logic

Recovery occurs when the system detects that a cached branch can no longer accept writes. This includes branches that have been deleted, merged, or are currently in the process of being deleted.

Writability Classification

The system identifies unwritable branches based on the status attribute returned by the Infrahub SDK.

Source Signal	Classification	System Action
`BranchNotFoundError`	Missing (Deleted)	Clear cache, recreate branch.
`BranchStatus.MERGED`	Unwritable (Read-only)	Clear cache, recreate branch, issue warning.
`BranchStatus.DELETING`	Unwritable	Clear cache, recreate branch, issue warning.
`BranchStatus.OPEN`	Writable	Reuse cached branch.

Sources: data-model.md:42-50, tasks.md:65-72

Proactive vs. Reactive Detection

The system employs two layers of detection to ensure reliability:

Proactive Detection: Before routing a write, get_or_create_session_branch inspects branch.get().status.
Reactive Net (FR-011): If a branch is merged during the write window (TOCTOU), the write tool catches the GraphQLError (read-only/merged), clears the session entry, and raises a retryable ToolError.

Sources: plan.md:27-33, spec.md:121-125

Session Branch Reset Tool

The reset_session_branch tool allows users to manually manage the active branch for their current session. It is tagged with write and session tags.

Tool Parameters and Responses

The tool accepts an optional branch argument to either clear the cache or switch to a specific target.

Parameter	Type	Description
`branch`	`str \| None`	The name of the branch to switch to. If `None`, the cache is cleared.

Action	Condition	Response Keys
Reset	`branch=None`	`action: "reset"`, `session_branch: null`
Switched	Existing writable branch	`action: "switched"`, `session_branch: <name>`
Created	Missing conformant branch	`action: "created"`, `created: true`, `session_branch: <name>`

Sources: spec.md:55-70, tasks.md:95-103

Branch Name Conformance

When switching to a non-existent branch, the name must match the configured branch_pattern. The system generates a regex using string.Formatter().parse(pattern) with the following placeholder mappings:

{date}: \d{8}
{hex}: [0-9a-f]{8}
{user}: [A-Za-z0-9._/-]+?

Sources: data-model.md:52-56, tasks.md:118-124

Security Hardening

The implementation includes protections to prevent unauthorized writes or privileged operations.

Default Branch Protection: The system rejects any attempt to switch the session to the instance's default branch (e.g., main). This is verified using the branch's is_default flag rather than a name comparison.
Mutation Filtering: mutate_graphql rejects branch-management mutations (e.g., BranchMerge, BranchDelete) and schema mutations to prevent bypassing the propose_changes review gate.
Isolation: A reset or recovery in one session does not affect the branch used by any other concurrent session.

Sources: tasks.md:144-155, spec.md:126-130

Implementation Details

The following Python snippet demonstrates the per-session object resolution used for dictionary keying:

# src/infrahub_mcp/utils.py

def _session_obj(ctx: Context) -> Any:
    rc = ctx.request_context
    if rc is None or getattr(rc, "session", None) is None:
        msg = "request_context.session must not be None"
        raise RuntimeError(msg)
    return rc.session

Sources: data-model.md:32-38, tasks.md:46-50

Session Branch Recovery and Reset addresses the wedge condition caused by merged branches. By scoping branch state to individual sessions and implementing proactive writability checks, the system ensures a continuous and isolated workflow for infrastructure mutations.