Spaces:
Running
Running
| 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()) | |