index.ts 5.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210
  1. import { BusEvent } from "@/bus/bus-event"
  2. import { Bus } from "@/bus"
  3. import z from "zod"
  4. import { Log } from "../util/log"
  5. import { Identifier } from "../id/id"
  6. import { Plugin } from "../plugin"
  7. import { Instance } from "../project/instance"
  8. import { Wildcard } from "../util/wildcard"
  9. export namespace Permission {
  10. const log = Log.create({ service: "permission" })
  11. function toKeys(pattern: Info["pattern"], type: string): string[] {
  12. return pattern === undefined ? [type] : Array.isArray(pattern) ? pattern : [pattern]
  13. }
  14. function covered(keys: string[], approved: Record<string, boolean>): boolean {
  15. const pats = Object.keys(approved)
  16. return keys.every((k) => pats.some((p) => Wildcard.match(k, p)))
  17. }
  18. export const Info = z
  19. .object({
  20. id: z.string(),
  21. type: z.string(),
  22. pattern: z.union([z.string(), z.array(z.string())]).optional(),
  23. sessionID: z.string(),
  24. messageID: z.string(),
  25. callID: z.string().optional(),
  26. message: z.string(),
  27. metadata: z.record(z.string(), z.any()),
  28. time: z.object({
  29. created: z.number(),
  30. }),
  31. })
  32. .meta({
  33. ref: "Permission",
  34. })
  35. export type Info = z.infer<typeof Info>
  36. export const Event = {
  37. Updated: BusEvent.define("permission.updated", Info),
  38. Replied: BusEvent.define(
  39. "permission.replied",
  40. z.object({
  41. sessionID: z.string(),
  42. permissionID: z.string(),
  43. response: z.string(),
  44. }),
  45. ),
  46. }
  47. const state = Instance.state(
  48. () => {
  49. const pending: {
  50. [sessionID: string]: {
  51. [permissionID: string]: {
  52. info: Info
  53. resolve: () => void
  54. reject: (e: any) => void
  55. }
  56. }
  57. } = {}
  58. const approved: {
  59. [sessionID: string]: {
  60. [permissionID: string]: boolean
  61. }
  62. } = {}
  63. return {
  64. pending,
  65. approved,
  66. }
  67. },
  68. async (state) => {
  69. for (const pending of Object.values(state.pending)) {
  70. for (const item of Object.values(pending)) {
  71. item.reject(new RejectedError(item.info.sessionID, item.info.id, item.info.callID, item.info.metadata))
  72. }
  73. }
  74. },
  75. )
  76. export function pending() {
  77. return state().pending
  78. }
  79. export function list() {
  80. const { pending } = state()
  81. const result: Info[] = []
  82. for (const items of Object.values(pending)) {
  83. for (const item of Object.values(items)) {
  84. result.push(item.info)
  85. }
  86. }
  87. return result.sort((a, b) => a.id.localeCompare(b.id))
  88. }
  89. export async function ask(input: {
  90. type: Info["type"]
  91. message: Info["message"]
  92. pattern?: Info["pattern"]
  93. callID?: Info["callID"]
  94. sessionID: Info["sessionID"]
  95. messageID: Info["messageID"]
  96. metadata: Info["metadata"]
  97. }) {
  98. const { pending, approved } = state()
  99. log.info("asking", {
  100. sessionID: input.sessionID,
  101. messageID: input.messageID,
  102. toolCallID: input.callID,
  103. pattern: input.pattern,
  104. })
  105. const approvedForSession = approved[input.sessionID] || {}
  106. const keys = toKeys(input.pattern, input.type)
  107. if (covered(keys, approvedForSession)) return
  108. const info: Info = {
  109. id: Identifier.ascending("permission"),
  110. type: input.type,
  111. pattern: input.pattern,
  112. sessionID: input.sessionID,
  113. messageID: input.messageID,
  114. callID: input.callID,
  115. message: input.message,
  116. metadata: input.metadata,
  117. time: {
  118. created: Date.now(),
  119. },
  120. }
  121. switch (
  122. await Plugin.trigger("permission.ask", info, {
  123. status: "ask",
  124. }).then((x) => x.status)
  125. ) {
  126. case "deny":
  127. throw new RejectedError(info.sessionID, info.id, info.callID, info.metadata)
  128. case "allow":
  129. return
  130. }
  131. pending[input.sessionID] = pending[input.sessionID] || {}
  132. return new Promise<void>((resolve, reject) => {
  133. pending[input.sessionID][info.id] = {
  134. info,
  135. resolve,
  136. reject,
  137. }
  138. Bus.publish(Event.Updated, info)
  139. })
  140. }
  141. export const Response = z.enum(["once", "always", "reject"])
  142. export type Response = z.infer<typeof Response>
  143. export function respond(input: { sessionID: Info["sessionID"]; permissionID: Info["id"]; response: Response }) {
  144. log.info("response", input)
  145. const { pending, approved } = state()
  146. const match = pending[input.sessionID]?.[input.permissionID]
  147. if (!match) return
  148. delete pending[input.sessionID][input.permissionID]
  149. Bus.publish(Event.Replied, {
  150. sessionID: input.sessionID,
  151. permissionID: input.permissionID,
  152. response: input.response,
  153. })
  154. if (input.response === "reject") {
  155. match.reject(new RejectedError(input.sessionID, input.permissionID, match.info.callID, match.info.metadata))
  156. return
  157. }
  158. match.resolve()
  159. if (input.response === "always") {
  160. approved[input.sessionID] = approved[input.sessionID] || {}
  161. const approveKeys = toKeys(match.info.pattern, match.info.type)
  162. for (const k of approveKeys) {
  163. approved[input.sessionID][k] = true
  164. }
  165. const items = pending[input.sessionID]
  166. if (!items) return
  167. for (const item of Object.values(items)) {
  168. const itemKeys = toKeys(item.info.pattern, item.info.type)
  169. if (covered(itemKeys, approved[input.sessionID])) {
  170. respond({
  171. sessionID: item.info.sessionID,
  172. permissionID: item.info.id,
  173. response: input.response,
  174. })
  175. }
  176. }
  177. }
  178. }
  179. export class RejectedError extends Error {
  180. constructor(
  181. public readonly sessionID: string,
  182. public readonly permissionID: string,
  183. public readonly toolCallID?: string,
  184. public readonly metadata?: Record<string, any>,
  185. public readonly reason?: string,
  186. ) {
  187. super(
  188. reason !== undefined
  189. ? reason
  190. : `The user rejected permission to use this specific tool call. You may try again with different parameters.`,
  191. )
  192. }
  193. }
  194. }