From b6c21339fd174275f5acfd61559e294a1ccf5f4b Mon Sep 17 00:00:00 2001
From: Donnie Adams <donnie@acorn.io>
Date: Fri, 31 Jan 2025 15:02:32 -0500
Subject: [PATCH] fix: modify Jira tool for API key support

The Jira API is different for OAuth versus API keys. This change
modifies the Jira tool so that it works with API keys and OAuth.

Signed-off-by: Donnie Adams <donnie@acorn.io>
---
 atlassian/credential/tool.gpt |  3 ++-
 atlassian/jira/index.js       | 31 ++++++++++++++++++++++++++-----
 atlassian/jira/tool.gpt       |  1 +
 3 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/atlassian/credential/tool.gpt b/atlassian/credential/tool.gpt
index 6e934d2a..2c15742d 100644
--- a/atlassian/credential/tool.gpt
+++ b/atlassian/credential/tool.gpt
@@ -2,7 +2,8 @@ Name: Atlassian OAuth Credential
 Share Credential: ../../oauth2 as atlassian
     with ATLASSIAN_OAUTH_TOKEN as token and
         atlassian as integration and
-        ATLASSIAN_OAUTH_TOKEN as prompt_tokens and
+        ATLASSIAN_API_TOKEN as prompt_tokens and
+        ATLASSIAN_EMAIL;ATLASSIAN_SITE_URL as prompt_vars and
         "offline_access
         read:me
         read:account
diff --git a/atlassian/jira/index.js b/atlassian/jira/index.js
index bfe516ba..0404ff2f 100644
--- a/atlassian/jira/index.js
+++ b/atlassian/jira/index.js
@@ -16,21 +16,38 @@ const command = process.argv[2]
 async function main() {
 
     try {
-        if (!process.env.ATLASSIAN_OAUTH_TOKEN) {
-            throw new Error("ATLASSIAN_OAUTH_TOKEN is required")
+        if (!process.env.ATLASSIAN_OAUTH_TOKEN && (!process.env.ATLASSIAN_API_TOKEN || !process.env.ATLASSIAN_EMAIL || !process.env.ATLASSIAN_SITE_URL)) {
+            throw new Error("ATLASSIAN_OAUTH_TOKEN or all of ATLASSIAN_API_TOKEN, ATLASSIAN_EMAIL, and ATLASSIAN_SITE_URL are required")
+        }
+
+        let authType = "Bearer"
+        let authToken = process.env.ATLASSIAN_OAUTH_TOKEN
+        if (process.env.ATLASSIAN_API_TOKEN) {
+            authType = "Basic"
+            authToken = Buffer.from(`${process.env.ATLASSIAN_EMAIL}:${process.env.ATLASSIAN_API_TOKEN.trim()}`).toString("base64")
         }
 
         let baseUrl
         if (command !== 'listJiraSites') {
-            if (!process.env.SITE_ID) {
+            if (process.env.ATLASSIAN_SITE_URL) {
+                baseUrl = process.env.ATLASSIAN_SITE_URL
+                if (baseUrl.endsWith('/')) {
+                    baseUrl = baseUrl.slice(0, -1)
+                }
+                if (!baseUrl.startsWith('http://') && !baseUrl.startsWith('https://')) {
+                    baseUrl = `https://${baseUrl}`
+                }
+                baseUrl = `${baseUrl}/rest/api/3`
+            } else if (process.env.SITE_ID) {
+                baseUrl = `https://api.atlassian.com/ex/jira/${process.env.SITE_ID}/rest/api/3`
+            } else {
                 throw new Error('site_id argument not provided')
             }
-            baseUrl = `https://api.atlassian.com/ex/jira/${process.env.SITE_ID}/rest/api/3`
         }
         const client = axios.create({
             baseURL: baseUrl,
             headers: {
-                'Authorization': `Bearer ${process.env.ATLASSIAN_OAUTH_TOKEN}`,
+                'Authorization': `${authType} ${authToken}`,
                 'Accept': 'application/json',
             },
         })
@@ -38,6 +55,10 @@ async function main() {
         let result = null
         switch (command) {
             case "listJiraSites":
+                if (process.env.ATLASSIAN_SITE_URL) {
+                    console.log("listJiraSites is not needed when ATLASSIAN_SITE_URL is set, so you can continue without needing the site_id argument")
+                    break
+                }
                 result = await listJiraSites(client)
                 break
             case "createIssue":
diff --git a/atlassian/jira/tool.gpt b/atlassian/jira/tool.gpt
index 21753c98..0427eea4 100644
--- a/atlassian/jira/tool.gpt
+++ b/atlassian/jira/tool.gpt
@@ -174,6 +174,7 @@ If you don't know what Jira Sites are available for the current user, call the L
 to get the list of sites that the user has access to.
 Tools that require a site_id must be called with a site_id from the Jira Sites that are available for the current user.
 When it's unclear which Jira Site a user is referring to, always ask the user to select a site before calling a tool that requires a site_id argument.
+When using the Jira Sites tool, it may tell you that you don't need a site_id. If it does, then that's great and you can move along without it.
 When information about the current user is needed to fulfill a request, call the Get Current User tool for the applicable Jira Sites.
 Always call Get Project to gather more information about a Project when prompted to interact with the issues in a project.
 Issue descriptions are always in ADF (Atlassian Document Format) and must be passed as a JSON string when calling the Create Issue tool.