Spaces:
Running
Running
File size: 11,486 Bytes
25a2c1c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | import asyncio,contextlib,hashlib,json,os,secrets,subprocess,sys,tempfile
from datetime import timedelta
from pathlib import Path
import httpx
from mcp import ClientSession,StdioServerParameters
from mcp.client.stdio import stdio_client
from mcp.shared.exceptions import McpError
from scripts.coverage_workflow import FIXED,STEPS
def result_data(result):
if result.structuredContent is not None:return result.structuredContent
text=[c.text for c in result.content if getattr(c,'type',None)=='text']
if len(text)==1:
try:
decoded=json.loads(text[0])
if isinstance(decoded,dict):return decoded
except ValueError:pass
return {'text':text}
async def main():
token=secrets.token_hex(24)
key=os.environ.get("NEEDLE_API_KEY", "test-local-key")
records=[];checks={};probes=[];native=[]
with tempfile.TemporaryDirectory() as td:
root=Path(td)
subprocess.run(['git','init'],cwd=root,capture_output=True,check=True)
(root/'app.py').write_text(FIXED)
subprocess.run(['git','add','app.py'],cwd=root,capture_output=True,check=True)
subprocess.run(['git','commit','-m','init'],cwd=root,capture_output=True,check=True)
def git(*args):return subprocess.run(['git',*args],cwd=root,capture_output=True,text=True,check=True).stdout.strip()
allowed=['read_file','write_file','run_tests','git_diff','git_commit','git_status']
def permissions(allowed_tools):
(root/'.git'/'needle-executor-policy.json').write_text(json.dumps({'allowed_tools':allowed_tools,'workspace':str(root)}))
permissions(allowed)
params=StdioServerParameters(command=sys.executable,args=['-m','scripts.workflow_mcp_server','--root',str(root)],env=dict(os.environ,NEEDLE_MCP_TEST_TOKEN=token,PYTHONPATH=str(Path.cwd()),PYTHONDONTWRITEBYTECODE='1'))
async with stdio_client(params) as (read,write):
async with ClientSession(read,write) as session:
await session.initialize()
tools=[t.model_dump(exclude_none=True) for t in (await session.list_tools()).tools if not t.name.startswith('probe_')]
checks['actual_mcp_tool_catalogue']=len(tools)==8
async def call(name,args,**kwargs):return await session.call_tool(name,args,meta={'capability':token},**kwargs)
unauth=await session.call_tool('read_file',{'path':'/project/app.py'},meta={'capability':'invalid'})
checks['mcp_authentication_enforced']=unauth.isError
checks['mcp_write_target_scope']=(await call('write_file',{'path':'/outside/app.py','content':'must not be written'})).isError
async with httpx.AsyncClient(base_url="http://127.0.0.1:7860", headers={'Authorization':'Bearer '+key}, timeout=30.0) as client:
resp = await client.get('/ready', headers={'Authorization':'Bearer invalid'})
checks['gateway_authentication_enforced']=resp.status_code==401
namespace='mcp-fixture'
(await client.post('/v1/tools/register',json={'namespace':namespace,'tools':tools})).raise_for_status()
version=(await client.get('/v1/tools/versions')).json()['versions'][namespace]
ref={'version':1,'namespace':namespace,'sha256':version};history=[];diff_only=False;tests_passed=False;write_exact=False
standard={'model':'needle3','tools':[{'type':'function','function':t} for t in tools],'messages':[{'role':'user','content':'git_status'}],'decision_context':{'workspace':'/project'}}
result=(await client.post('/v1/chat/completions',json=standard)).json()
standard_calls=result['choices'][0]['message'].get('tool_calls') or []
ok=len(standard_calls)==1 and standard_calls[0]['function']['name']=='git_status' and json.loads(standard_calls[0]['function']['arguments'])=={} and result['needle']['outcome']=='selected'
checks['standard_openai_tool_format']=ok and not (await call('git_status',{})).isError
probes.append({'name':'standard-openai-tools','request':standard,'response':result})
for index,(query,name,args,deny) in enumerate(STEPS):
policy={'workspace':'/project','session_id':'mcp-workflow','client_id':'official-mcp','task_id':str(index),'arguments':{name:args} if name and query!='Open the file we identified.' else {},'denied_tools':['run_tests'] if deny else []}
body={'model':'needle3','registry_reference':ref,'messages':history[-2:]+[{'role':'user','content':query}],'decision_context':policy}
native_start=len(native)
r=await client.post('/v1/chat/completions',json=body);r.raise_for_status();data=r.json();calls=data['choices'][0]['message'].get('tool_calls') or []
parsed=[{'name':c['function']['name'],'arguments':json.loads(c['function']['arguments'])} for c in calls]
correct=(parsed==[{'name':name,'arguments':args}] and data['needle']['outcome']=='selected') if name else not parsed and data['needle']['outcome'] in {'abstain','confirm'}
output=None
if name and correct:
if name=='write_file':
original_bytes=(root/'app.py').read_bytes()
permissions([n for n in allowed if n!='write_file'])
denied=await call(name,args)
checks['executor_rechecks_permission_after_selection']=denied.isError and (root/'app.py').read_bytes()==original_bytes
permissions(allowed)
result_val=await call(name,args);output=result_val.model_dump(mode='json');output['decoded_fixture_data']=result_data(result_val)
assert not result_val.isError,(index,result_val)
structured=result_data(result_val)
if name=='run_tests' and structured.get('returncode')==0:tests_passed=True
if name=='write_file':write_exact=(root/'app.py').read_bytes()==FIXED.encode()
if name=='git_diff':diff_only=git('diff','--name-only').splitlines()==['app.py']
records.append({'step':index,'request':body,'response':data,'exact_route':correct,'mcp_result':output,'native':native[native_start:]})
history.append({'role':'user','content':query})
if output is not None:
history.append({'role':'tool','content':json.dumps(output['decoded_fixture_data'])})
checks.update(all_routes_exact=all(r['exact_route'] for r in records),file_bytes_exact=write_exact,tests_passed_after_edit=tests_passed,executor_revocation_probe_ran=checks.get('executor_rechecks_permission_after_selection',False),diff_only_expected_file=diff_only,commit_only_expected_file=git('show','--format=','--name-only','HEAD').splitlines()==['app.py'],commit_message_exact=git('log','-1','--format=%s')=='Repair addition',repository_clean=not git('status','--porcelain'))
revoked_content=FIXED+'# must not execute after revocation\n'
request={'model':'needle3','tools':tools,'messages':[{'role':'user','content':'write_file /project/app.py content '+json.dumps(revoked_content)}],'decision_context':{'workspace':'/project','arguments':{'write_file':{'path':'/project/app.py','content':revoked_content}}}}
request['decision_context']['denied_tools']=['write_file']
denied_route=(await client.post('/v1/chat/completions',json=request)).json();checks['gateway_current_permission_denial']=not denied_route['choices'][0]['message']['tool_calls']
current={'model':'needle3','registry_reference':ref,'messages':[{'role':'user','content':'git_status'}],'decision_context':{'workspace':'/project','session_id':'mcp-warm','client_id':'official-mcp','task_id':'warm'}}
full=(await client.post('/v1/chat/completions',json=current)).json();revision=full['needle']['revision'];outcomes=[]
for _ in range(3):
current['decision_context'].update(incremental=True,previous_revision=revision)
d=(await client.post('/v1/chat/completions',json=current)).json();revision=d['needle']['revision'];outcomes.append(d['needle']['outcome'])
if d['needle']['outcome']=='selected':assert not (await call('git_status',{})).isError
checks['incremental_native_and_mcp_calls']=outcomes==['selected']*3
updated=[{**t,'description':t['description']+' Updated catalogue.'} for t in tools]
(await client.post('/v1/tools/register',json={'namespace':namespace,'tools':updated})).raise_for_status()
checks['stale_registry_rejected']=(await client.post('/v1/chat/completions',json=current)).status_code==409
updated_sha=(await client.get('/v1/tools/versions')).json()['versions'][namespace]
fresh={'model':'needle3','registry_reference':{**ref,'sha256':updated_sha},'messages':[{'role':'user','content':'git_status'}],'decision_context':{'workspace':'/project','session_id':'mcp-updated-registry'}}
result=(await client.post('/v1/chat/completions',json=fresh)).json();calls=result['choices'][0]['message'].get('tool_calls') or []
ok=len(calls)==1 and calls[0]['function']['name']=='git_status' and json.loads(calls[0]['function']['arguments'])=={} and result['needle']['outcome']=='selected'
checks['updated_registry_routes_and_executes']=ok and not (await call('git_status',{})).isError
probes.append({'name':'updated-registry','request':fresh,'response':result})
checks['mcp_error_propagation']=(await call('probe_error',{})).isError
timeout=False
try:await call('probe_delay',{},read_timeout_seconds=timedelta(seconds=.05))
except McpError as exc:timeout=exc.error.code==408
checks['mcp_timeout']=timeout
task=asyncio.create_task(call('probe_delay',{}));await asyncio.sleep(.05);task.cancel()
with contextlib.suppress(asyncio.CancelledError):await task
checks['mcp_cancellation']=task.cancelled()
checks['mcp_usable_after_cancel']=not (await call('read_file',{'path':'/project/app.py'})).isError
result={'methodology':'Real official mcp==1.26.0 stdio ClientSession/FastMCP transport against external constrained Needle 3.0.1 Docker gateway. Test capabilities passed only in MCP metadata.','checks':checks,'success':all(checks.values()),'source_sha256':hashlib.sha256((root/'app.py').read_bytes()).hexdigest(),'records':records,'probes':probes}
Path('benchmarks/workflow/results/mcp-integration-docker.json').write_text(json.dumps(result,indent=2)+'\n')
print(json.dumps({'success':result['success'],'checks':checks},indent=2))
if not result['success']:raise SystemExit('Real MCP qualification failed; evidence retained')
if __name__=='__main__':asyncio.run(main())
|