fix: refactor mcp with visitor

[ravenac95]

Bit of a major refactor to try to make this code more maintainable and correct. I have found some additional bugs since making this but I think we need to review this now. Otherwise it'll be too large a change for anyone to look at.

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
kariba-network	Ready	Preview, Comment	Feb 5, 2026 5:01pm
oso-www	Ready	Preview, Comment	Feb 5, 2026 5:01pm

[github-actions]

Code Review - High Signal Issues Found

Issue 1: Missing discriminator field for union types

File: warehouse/oso_mcp/oso_mcp/server/graphql/pydantic_generator.py
Method: handle_leave_union (around line 1210-1217 in the diff)
Description: When adding a union field to the parent context, the discriminator field is not being set.
Reason: Logic error - union fields require discriminator="typename__" in Pydantic for proper deserialization

The code currently does:

if not is_required:
    python_type = t.Optional[python_type]
    field_info = Field(default=None)
else:
    field_info = Field()

current_context = self.require_model_context("add union field")

But it should include the discriminator:

if not is_required:
    python_type = t.Optional[python_type]
    field_info = Field(default=None, discriminator="typename__")
else:
    field_info = Field(discriminator="typename__")

current_context = self.require_model_context("add union field")

This will cause runtime errors when trying to deserialize GraphQL responses containing union types because Pydantic won't know how to discriminate between the different union member types.

[github-actions]

Bug: The assertion error message incorrectly references schema_start (which is a GraphQLObjectType object) instead of fragment.type_condition_name (which is the string name). This will cause a poorly formatted error message if the assertion fails.

The error message should be:

Suggested change

	f"Schema does not have a {schema_start} type defined"
	f"Schema does not have a {fragment.type_condition_name} type defined"

[github-actions]

Missing discriminator parameter for union field. Pydantic requires discriminator="typename__" to deserialize union types correctly based on the __typename field from GraphQL responses. This will cause runtime failures when deserializing responses containing union types. See lines 584-586 where the discriminator is correctly set for comparison.

[github-actions]

Non-exhaustive pattern matching - if current_context is neither PydanticModelBuildContext nor UnionTypeBuildContext, the match silently does nothing. Add a default case to raise an error or use require_model_context() to ensure type safety and prevent silent data loss.