This should probably come from the current thread; perhaps it should just be keyed on PsIsSystemThread(KeGetCurrentThread()). Ideally we'd have at least some tests.
I agree that that's about the best we can do. It won't work for Zw*() calls though, since we just forward Zw*() to Nt*() counterparts just like ntdll.