fix: propagate context to MCP server clean

[CakeCrusher]

TLDR:
The context passed by the InstrumentedStreamWriter.send was being mutated by processes preceeding InstrumentedStreamReader resulting in disjointed MCP traces.

By adding instrumentation at _handle_request (an earlier state in the program flow) the original message composed by InstrumentedStreamWriter.send is preserved. This context is then propagated at the server level.

the context is detached by the same method it was being detached before, through InstrumentedStreamReader.

[CakeCrusher]

Hi @mikeldking heres the new MCP propagation PR in case i lost you in my mess

[mikeldking]

Seems to be failing tests. Can you check the failures?

[mikeldking]

Suggested change

	__version__ = "1.3.1"
	__version__ = "1.3.0"

[CakeCrusher]

Does CD handle it?

[mikeldking]

This doesn't match the method instrumented above?

[CakeCrusher]

@mikeldking
The problem with the tests seems to be that the socket network restrictions that exist in the CI actions environment is provoking the test_hello inside test_instrumentor to stall.
For the time being I removed testing on test_hello inside test_instrumentor.
I can think of 3 better solutions:

1. Add a github env variable that you can use to conditionally run transport options sse and streamable-http. This enables complete testing on local but limited testing on CICD runner environment.
2. Resolve the network limits only on this MCP.
3. Edit restrictions on CICD liklely related to pytest-socket plugin.

[graphite-app]

The current try/finally structure could lead to unexpected behavior. If an exception occurs during context extraction, the wrapped function would still execute in the finally block, which is likely not the intended behavior.

Consider restructuring like this:

token = None
try:
    # Message has been deserialized, we need to extract the traceparent
    _meta = {"traceparent": args[1].params.meta.traceparent}
    ctx = propagate.extract(_meta)
    token = context.attach(ctx)
    
    # Execute the wrapped function inside the try block
    res = await wrapped(*args, **kwargs)
    return res
finally:
    # Only detach if we successfully attached
    if token:
        context.detach(token)

This ensures the wrapped function only executes if context extraction succeeds, while still guaranteeing context detachment in all cases where it was attached.

Suggested change

	token = None
	try:
	# Message has been deserialized, we need to extract the traceparent
	_meta = {"traceparent": args[1].params.meta.traceparent}
	ctx = propagate.extract(_meta)
	token = context.attach(ctx)
	finally:
	res = await wrapped(*args, **kwargs)
	if token:
	context.detach(token)
	return res
	token = None
	try:
	# Message has been deserialized, we need to extract the traceparent
	_meta = {"traceparent": args[1].params.meta.traceparent}
	ctx = propagate.extract(_meta)
	token = context.attach(ctx)
	# Execute the wrapped function inside the try block
	res = await wrapped(*args, **kwargs)
	return res
	finally:
	# Only detach if we successfully attached
	if token:
	context.detach(token)

Spotted by Diamond

Is this helpful? React 👍 or 👎 to let us know.