Skip to content

[Official SDK Migration] Migrate session and metrics hooks - #511

Open
Shweta-Deshpande wants to merge 4 commits into
mainfrom
TFECO-13096
Open

[Official SDK Migration] Migrate session and metrics hooks#511
Shweta-Deshpande wants to merge 4 commits into
mainfrom
TFECO-13096

Conversation

@Shweta-Deshpande

Copy link
Copy Markdown
Collaborator

What changed

Adds the last piece of the official go-sdk migration: session hooks and metrics hooks, replacing the mark3labs server.Hooks equivalents for the official SDK path.

  • pkg/mcp-official/server.go: Added WithOnSession, a new Option that runs once per session, right after the client finishes initializing. This is the go-sdk's only per-session hook point.
  • cmd/terraform-mcp-server/init.go: Wired WithOnSession into getOfficialStreamableServer to pre-create the session's TFE/HTTP clients on connect, then wait in the background and clean them up (plus rate-limit state) on disconnect. This replaces the old AddOnRegisterSession/AddOnUnregisterSession hooks.
  • pkg/mcp-official/tools/middleware/metrics.go (new): A middleware that records tool-call counts, latency, errors, and client type — the go-sdk equivalent of the mark3labs AddAfterInitialize/AddBeforeCallTool/AddAfterCallTool hooks used for metrics. It reads client info straight from the session instead of needing a separate cache.
  • pkg/client/metrics.go: Split RecordToolCall into a common RecordToolCallByName, so both the mark3labs hooks and the new go-sdk middleware share the same recording logic.
  • Added tests for both the new metrics middleware and the session cleanup lifecycle (session_lifecycle_test.go), the connecting and disconnecting flow, that were previously not covered.

PCI review checklist

  • I have documented a clear reason for, and description of, the change I am making.

  • If applicable, I've documented a plan to revert these changes if they require more than reverting the pull request.

  • If applicable, I've documented the impact of any changes to security controls.

    Examples of changes to security controls include using new access control methods, adding or removing logging pipelines, etc.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant