next.ts 7.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269
  1. import { Bus } from "@/bus"
  2. import { BusEvent } from "@/bus/bus-event"
  3. import { Config } from "@/config/config"
  4. import { Identifier } from "@/id/id"
  5. import { Instance } from "@/project/instance"
  6. import { Storage } from "@/storage/storage"
  7. import { fn } from "@/util/fn"
  8. import { Log } from "@/util/log"
  9. import { Wildcard } from "@/util/wildcard"
  10. import z from "zod"
  11. export namespace PermissionNext {
  12. const log = Log.create({ service: "permission" })
  13. export const Action = z.enum(["allow", "deny", "ask"]).meta({
  14. ref: "PermissionAction",
  15. })
  16. export type Action = z.infer<typeof Action>
  17. export const Rule = z
  18. .object({
  19. permission: z.string(),
  20. pattern: z.string(),
  21. action: Action,
  22. })
  23. .meta({
  24. ref: "PermissionRule",
  25. })
  26. export type Rule = z.infer<typeof Rule>
  27. export const Ruleset = Rule.array().meta({
  28. ref: "PermissionRuleset",
  29. })
  30. export type Ruleset = z.infer<typeof Ruleset>
  31. export function fromConfig(permission: Config.Permission) {
  32. const ruleset: Ruleset = []
  33. for (const [key, value] of Object.entries(permission)) {
  34. if (typeof value === "string") {
  35. ruleset.push({
  36. permission: key,
  37. action: value,
  38. pattern: "*",
  39. })
  40. continue
  41. }
  42. ruleset.push(...Object.entries(value).map(([pattern, action]) => ({ permission: key, pattern, action })))
  43. }
  44. return ruleset
  45. }
  46. export function merge(...rulesets: Ruleset[]): Ruleset {
  47. return rulesets.flat()
  48. }
  49. export const Request = z
  50. .object({
  51. id: Identifier.schema("permission"),
  52. sessionID: Identifier.schema("session"),
  53. permission: z.string(),
  54. patterns: z.string().array(),
  55. metadata: z.record(z.string(), z.any()),
  56. always: z.string().array(),
  57. tool: z
  58. .object({
  59. messageID: z.string(),
  60. callID: z.string(),
  61. })
  62. .optional(),
  63. })
  64. .meta({
  65. ref: "PermissionRequest",
  66. })
  67. export type Request = z.infer<typeof Request>
  68. export const Reply = z.enum(["once", "always", "reject"])
  69. export type Reply = z.infer<typeof Reply>
  70. export const Approval = z.object({
  71. projectID: z.string(),
  72. patterns: z.string().array(),
  73. })
  74. export const Event = {
  75. Asked: BusEvent.define("permission.asked", Request),
  76. Replied: BusEvent.define(
  77. "permission.replied",
  78. z.object({
  79. sessionID: z.string(),
  80. requestID: z.string(),
  81. reply: Reply,
  82. }),
  83. ),
  84. }
  85. const state = Instance.state(async () => {
  86. const projectID = Instance.project.id
  87. const stored = await Storage.read<Ruleset>(["permission", projectID]).catch(() => [] as Ruleset)
  88. const pending: Record<
  89. string,
  90. {
  91. info: Request
  92. resolve: () => void
  93. reject: (e: any) => void
  94. }
  95. > = {}
  96. return {
  97. pending,
  98. approved: stored,
  99. }
  100. })
  101. export const ask = fn(
  102. Request.partial({ id: true }).extend({
  103. ruleset: Ruleset,
  104. }),
  105. async (input) => {
  106. const s = await state()
  107. const { ruleset, ...request } = input
  108. for (const pattern of request.patterns ?? []) {
  109. const rule = evaluate(request.permission, pattern, ruleset, s.approved)
  110. log.info("evaluated", { permission: request.permission, pattern, action: rule })
  111. if (rule.action === "deny")
  112. throw new DeniedError(ruleset.filter((r) => Wildcard.match(request.permission, r.permission)))
  113. if (rule.action === "ask") {
  114. const id = input.id ?? Identifier.ascending("permission")
  115. return new Promise<void>((resolve, reject) => {
  116. const info: Request = {
  117. id,
  118. ...request,
  119. }
  120. s.pending[id] = {
  121. info,
  122. resolve,
  123. reject,
  124. }
  125. Bus.publish(Event.Asked, info)
  126. })
  127. }
  128. if (rule.action === "allow") continue
  129. }
  130. },
  131. )
  132. export const reply = fn(
  133. z.object({
  134. requestID: Identifier.schema("permission"),
  135. reply: Reply,
  136. message: z.string().optional(),
  137. }),
  138. async (input) => {
  139. const s = await state()
  140. const existing = s.pending[input.requestID]
  141. if (!existing) return
  142. delete s.pending[input.requestID]
  143. Bus.publish(Event.Replied, {
  144. sessionID: existing.info.sessionID,
  145. requestID: existing.info.id,
  146. reply: input.reply,
  147. })
  148. if (input.reply === "reject") {
  149. existing.reject(input.message ? new CorrectedError(input.message) : new RejectedError())
  150. // Reject all other pending permissions for this session
  151. const sessionID = existing.info.sessionID
  152. for (const [id, pending] of Object.entries(s.pending)) {
  153. if (pending.info.sessionID === sessionID) {
  154. delete s.pending[id]
  155. Bus.publish(Event.Replied, {
  156. sessionID: pending.info.sessionID,
  157. requestID: pending.info.id,
  158. reply: "reject",
  159. })
  160. pending.reject(new RejectedError())
  161. }
  162. }
  163. return
  164. }
  165. if (input.reply === "once") {
  166. existing.resolve()
  167. return
  168. }
  169. if (input.reply === "always") {
  170. for (const pattern of existing.info.always) {
  171. s.approved.push({
  172. permission: existing.info.permission,
  173. pattern,
  174. action: "allow",
  175. })
  176. }
  177. existing.resolve()
  178. const sessionID = existing.info.sessionID
  179. for (const [id, pending] of Object.entries(s.pending)) {
  180. if (pending.info.sessionID !== sessionID) continue
  181. const ok = pending.info.patterns.every(
  182. (pattern) => evaluate(pending.info.permission, pattern, s.approved).action === "allow",
  183. )
  184. if (!ok) continue
  185. delete s.pending[id]
  186. Bus.publish(Event.Replied, {
  187. sessionID: pending.info.sessionID,
  188. requestID: pending.info.id,
  189. reply: "always",
  190. })
  191. pending.resolve()
  192. }
  193. // TODO: we don't save the permission ruleset to disk yet until there's
  194. // UI to manage it
  195. // await Storage.write(["permission", Instance.project.id], s.approved)
  196. return
  197. }
  198. },
  199. )
  200. export function evaluate(permission: string, pattern: string, ...rulesets: Ruleset[]): Rule {
  201. const merged = merge(...rulesets)
  202. log.info("evaluate", { permission, pattern, ruleset: merged })
  203. const match = merged.findLast(
  204. (rule) => Wildcard.match(permission, rule.permission) && Wildcard.match(pattern, rule.pattern),
  205. )
  206. return match ?? { action: "ask", permission, pattern: "*" }
  207. }
  208. const EDIT_TOOLS = ["edit", "write", "patch", "multiedit"]
  209. export function disabled(tools: string[], ruleset: Ruleset): Set<string> {
  210. const result = new Set<string>()
  211. for (const tool of tools) {
  212. const permission = EDIT_TOOLS.includes(tool) ? "edit" : tool
  213. const rule = ruleset.findLast((r) => Wildcard.match(permission, r.permission))
  214. if (!rule) continue
  215. if (rule.pattern === "*" && rule.action === "deny") result.add(tool)
  216. }
  217. return result
  218. }
  219. /** User rejected without message - halts execution */
  220. export class RejectedError extends Error {
  221. constructor() {
  222. super(`The user rejected permission to use this specific tool call.`)
  223. }
  224. }
  225. /** User rejected with message - continues with guidance */
  226. export class CorrectedError extends Error {
  227. constructor(message: string) {
  228. super(`The user rejected permission to use this specific tool call with the following feedback: ${message}`)
  229. }
  230. }
  231. /** Auto-rejected by config rule - halts execution */
  232. export class DeniedError extends Error {
  233. constructor(public readonly ruleset: Ruleset) {
  234. super(
  235. `The user has specified a rule which prevents you from using this specific tool call. Here are some of the relevant rules ${JSON.stringify(ruleset)}`,
  236. )
  237. }
  238. }
  239. export async function list() {
  240. return state().then((x) => Object.values(x.pending).map((x) => x.info))
  241. }
  242. }